けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

yukicoder No.843 Triple Primes

エラトステネスの篩の練習問題に良さそう

問題概要

 1 以上  N 以下の素数の組  (p, q, r) であって、

  •  p + q = r^{2}

を満たすものの個数を求めよ。

制約

  •  1 \le N \le 5 \times 10^{5}

考えたこと

 N = 1 のときは明らかにダメ。

また、偶数かつ素数である整数は 2 のみであることに注意する。 p, q がともに奇数とすると  r は偶数となる。よって  r = 2 となるが、 r = 2 を満たす解は  (p, q, r) = (2, 2, 2) しかない。よって  p, q がともに奇数であることはありえない。

以上から、 p = 2 または  q = 2 であるとしてよい。ここでは  q = 2 として解いて最後に 2 倍することにする。 p = q = 2 の場合をあらかじめ特別扱いしておく。

 p + 2 = r^{2}

を満たす奇素数  (p, r) をすべて求める問題へと帰着された。次のように解ける。

  • あらかじめエラトステネスの篩によって  N 以下のすべての整数について、素数かどうかを求めておく ( O(N \log\log N))
  •  r \times r - 2 \le N を満たすすべての素数  r に対して  r^{2} - 2素数ならば、答えに加算する

全体で計算量は  O(N \log\log N) となる。

コード

#include <bits/stdc++.h>
using namespace std;

vector<bool> isprime;
vector<int> Era(int n) {
    isprime.resize(n, true);
    vector<int> res;
    isprime[0] = false; isprime[1] = false;
    for (int i = 2; i < n; ++i) isprime[i] = true;
    for (int i = 2; i < n; ++i) {
        if (isprime[i]) {
            res.push_back(i);
            for (int j = i*2; j < n; j += i) isprime[j] = false;
        }
    }
    return res;
}

long long solve(long long N) {
    Era(1000000);
    if (N == 1) return 0;
    long long res = 1; // (2, 2, 2)
    for (long long r = 3; r*r - 2 <= N; ++r) {
        if (isprime[r] && isprime[r*r-2]) res += 2;
    }
    return res;
}

int main() {
    int N;
    cin >> N;
    cout << solve(N) << endl;
}