公式

C - 2026 解説 by Nyaan


まず、いったん計算量を無視した解法を考えてみましょう。\(N\) 以下の各整数が良い整数であるかどうかを判定できればよいので、例えば次のアルゴリズムが考えられます。

  • 答えを管理する配列 \(a\) を用意する。はじめ、\(a\) は空である。
  • \(n=1,2,\dots, N\) に対して次の操作を行う。
    • はじめ、\(c = 0\) とする。
    • \(0 \lt x \leq \sqrt{n}\) を満たす整数 \(x\) について次の操作を行う。(注:\(x \leq \sqrt{n}\) という条件を加えているのは、\(x \gt \sqrt{n}\) の場合 \(x^2+y^2=n\) を満たす \(y\) が存在しないためである。)
      • \(x^2 + y^2 = n\) を満たす非負の数 \(y\) を計算する。(これは sqrt(n - x * x) を計算すればよい。)
      • \(y\)\(x \lt y\) を満たす正整数である場合、\(c\) をインクリメントする。
    • 最終的な \(c\) の値が \(1\) である場合、良い整数としての条件を満たすので、\(a\)\(n\) を追加する。
  • 最終的な \(a\) が答えを格納した配列となる。

しかし、このアルゴリズムは \(\mathrm{O}(N \sqrt{N})\)\(y\) を計算しているため実装しても TLE してしまうので工夫が必要です。

上記のアルゴリズムにおいて調べている \((x, y)\) の範囲が \(x^2 + y^2 \leq N\) である点に注目します。ここから、はじめに \(n\) を固定することなく \((x, y)\) を調べていくという方法が導出できます。つまり、次のアルゴリズムが従います。

  • はじめに \(c = (c_1, c_2, \dots, c_N)\)\(0\) 初期化された配列とする。
  • \(x = 1, 2, \dots\) について \(x^2 \leq N\) を満たす間、次の操作を行う。
    • \(y = x+1, x+2, \dots\) について \(x^2 + y^2 \leq N\) を満たす間、次の操作を行う。
      • \(c_{x^2 + y^2}\) をインクリメントする。
  • 最終的に \(c_n = 1\) であるような \(n\) が良い整数となるから、\(c_n=1\) である \(n\) を列挙する。

このアルゴリズムの計算量を考えてみると、\(x, y\)\(x^2 + y^2 \leq N\) の範囲のみを探索していることから \(x, y \leq \sqrt{N}\) が従うので、調べる \((x,y)\) の個数は \(\mathrm{O}(\sqrt{N} \times \sqrt{N}) = \mathrm{O}(N)\) 個です。よってこの問題を \(\mathrm{O}(N)\) で解くことが出来ていて、これは十分高速です。

  • 実装例(C++)
#include <iostream>
#include <vector>
using namespace std;

int main() {
  int N;
  cin >> N;
  vector<int> c(N + 1);
  for (int x = 1; x * x <= N; x++) {
    for (int y = x + 1; x * x + y * y <= N; y++) {
      c[x * x + y * y]++;
    }
  }
  vector<int> ans;
  for (int n = 1; n <= N; n++) {
    if (c[n] == 1) ans.push_back(n);
  }
  cout << ans.size() << "\n";
  for (int i = 0; i < (int)ans.size(); i++) {
    cout << ans[i] << " \n]"[i + 1 == (int)ans.size()];
  }
}

投稿日時:
最終更新: