はまやんはまやんはまやん

hamayanhamayan's blog

3PrimeCounting [yukicoder No.732]

https://yukicoder.me/problems/no/732

解法

https://yukicoder.me/submissions/283736

cとabc=a+b+cをそれぞれ全探索する。
「N以下の素数の個数はN/logN個くらい」というのがあるため、それぞれ全探索でO(N^2/log^2N)。
O(N^2)は厳しいが、これなら確かに見た感じ行けそう。
cはN以下の素数、abcは3N以下の素数で全探索する。
 
あとは、c,abcが固定のときに条件を満たすa,bの組合せをO(1)で求めていきたい。
そのために、「cnt[i] := c未満の2つの素数を足してiとなる素数の組の組合せ」を更新しながら使っていく。
c,abcが分かっている時、cnt[abc-c]が条件を満たす組み合わせ数となるので、それを足し合わせていく。
cnt配列の問題がc未満という部分であるが、これは更新しながら使うことで実現できる。
 
あるcでのループを考える。(次のcをc'とする)
cnt[i]がc未満に対して正しく作られているとする。
とりあえず、cnt[abc-c]を使って答えを足していく。
次にcnt[i]を「c未満に対して」から「c'未満に対して」に変換する。
「c'未満に対して」は言い換えると「c以下に対して」とも言える。
なので、bがc, aがc未満の素数である場合のcntを追加してやれば、「c以下に対して」が実現できる。
よって、c未満の素数をvaで保持しておくと、vaの1つをa, b=cとして、cntを更新できる。
vaのループも高々O(N/logN)なので、全体の計算量は変わらず、O(N^2/log^2N)である。
 
なお、この問題はFFTによる解法が存在する。

int N;
ll cnt[301010];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N;

    auto vp = makePrimes(N);
    auto vp2 = makePrimes(N * 3);

    ll ans = 0;
    vector<int> va;
    fore(c, vp) {
        fore(abc, vp2) if (c < abc) ans += cnt[abc - c];
        fore(a, va) cnt[a + c]++;
        va.push_back(c);
    }
    cout << ans << endl;
}