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

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

yukicoder No.843 Triple Primes

yukicoder 数学(整数問題) 素数 バケット 集計処理 エラトステネスの篩 数え上げ問題 前処理 【問題集】整数変数の式で表された条件を扱う探索

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

問題概要

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;
}