2785:4 Values whose Sum is 0

keyword

定数最適化 C++

概要

数の集合A,B,C,Dが与えられる。それぞれの濃度はn(<4000)である。このとき、a+b+c+d==0となる組み合わせは何通りあるか求める問題。
この問題はアルゴリズムはずっと前から分かっていたけどTLEで弾かれつづけていた。固定長配列使うとかmapを使わないとか色々手を加えてやっと通った。
a+b == -(c+d)であることに気が付けばオーダーをO(n^2 log n)に落とせる。-(c+d)を先に計算、ソートしておいて全ての(a+b)に対して見つかった分だけカウントすれば良い。STLのcountは線形探索なのでequal_rangeとdistanceを使うのが良い。

int main(){
    int n, i, j, m;
    scanf("%d",&n);
    m = n*n;
    static int as[4001], bs[4001], cs[4001], ds[4001], abs[16000002];
    REP(i,n){
        scanf("%d%d%d%d", as+i, ds+i, bs+i, cs+i);
    }
    REP(i,n)REP(j,n){
        abs[i*n+j] = -as[i]-bs[j];
    }
    ll ans=0;
    sort(abs, abs+m);
    __typeof(equal_range(as,as+n,0)) range;
    REP(i,n)REP(j,n){
        range = equal_range(abs, abs+m, cs[i]+ds[j]);
        ans += (ll)distance(range.first, range.second);
    }
    printf("%lld\n", ans);
    return 0;
}