haraduka's diary

やる気が欲しい

Codeforces Round#314(Div2) C.Geometric Progression

数列が与えられて、その中にある公比の3項の等比数列が何通り存在するか、という問題。
まぁmapかなんかで適当に管理して、数列の1項目,2項目,3項目が何通り存在するかを順に更新していけば一瞬で解ける。

本番pretest7が全く通らなかった…オーバーフローだった…
AtCoderとかだと数が大きくなるためintでは間に合わない的なことが書いてあるけど、そういやCodeforcesって書いてないんだっけ…5,6ヶ月ぶりにやったらうまーくハマってしまった…

    int N, K; cin >> N >> K;
    int A[N];
    for(int i=0; i<N; i++){
        cin >> A[i];
    }
    map<int, array<ll, 3> > mp; //arrayの0,1,2がそれぞれ、1,2,3項目として何通り存在するか。
    for(int i=0; i<N; i++){
        int a = A[i];
        auto& mps = mp[a];
        if(a%K == 0 && mp.find(a/K) != mp.end()){
            auto& mpi = mp[a/K];
            ll a0 = mps[1]+mpi[0];
            ll a1 = mps[2]+mpi[1];
            mps[1] = a0;
            mps[2] = a1;
        }
        mps[0] += 1;
    }
    ll res = 0;
    for(auto it=mp.begin(); it!=mp.end(); it++){
        res += (it->second)[2];
    }
    cout << res << endl;