[yukicoder] No. 843 Triple Primes

問題

方針

素数の列挙

エラトステネスの篩を用いて素数を列挙します。

制約を考える

\( p + q = r^2 \) を満たすので、\( p, q, r \) のうち、二つの変数を探索すればいいことが分かります。ここで、\( 2 \leq p, q, r \leq N \) です。また、方程式から、\( r \) について、

\[
\begin{eqnarray}
2 &\leq& p \leq N\\
2 &\leq& r^2 – q \leq N\\
2 + q &\leq& r^2 \leq N + q\\
\sqrt{2 + q} &\leq& r \leq \sqrt{N + q} \ (r > 1)
\end{eqnarray}
\]

を満たす必要があります。したがって、\( q, r \) について全探索を行います。

コード

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

static const int MAX = 500005;
bool isPrime[MAX];

void aryPrime() {
  for (int i = 0; i < MAX; i++) isPrime[i] = true;
  isPrime[0] = false;
  isPrime[1] = false;
  int n = sqrt(MAX);
  for (int i = 2; i <= n; i++) {
    if (!isPrime[i]) continue;
    for (int j = i * 2; j < MAX; j += i) {
      isPrime[j] = false;
    }
  }
}
int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int N;
  cin >> N;
  vector<ll> p;
  aryPrime();
  for (int i = 2; i <= N; i++) {
    if (isPrime[i]) {
      p.push_back(i);
    }
  }
  
  int ans = 0;

  for (int i = 0; i < p.size(); i++) {
    for (int j = 0; j < p.size(); j++) {
      ll t = p[i] * p[i] - p[j];
      if (t < 1 || t > N) break;
      if (isPrime[t]) {
        ans++;
      }
    }
  }
  cout << ans << "\n";
  return 0;
}