次の問題を考えます。

\(N\) 以下の素数 \(p\) 全てに対して、\(a^2+b^2=p\) を満たす正整数の組 \((a,b)\) を求めよ。

ナイーブな全探索により \(O(N^{1.5})\) で求めることができます。以下では \(O(N\log N)\) で求める方法を説明します。

\(O(N\log N)\) の前計算により、 \(N\) 以下の数の素因数分解は \(O(\log N)\) でできるとしてよいです。（参考）

\(p\) の昇順に計算します。

まず、\(x^2+1\equiv 0 \mod p\) となる \(x\) を以下の手順で見つけます。

• 整数 \(y\) をランダムに取り、\(x=y^{(p-1)/4} \bmod p\) とする。\(x^2+1\equiv 0 \mod p\) でなければ \(y\) を選び直す。

\(y\) をランダムに取るとき、目的のものが得られる確率は \(1/2\) なので、この操作は \(O(1)\) 回のやり直しで済みます。

この \(x\) を用いて \(x^2+1=kp\) とすると \(k<p\) なので \(k\) のガウス整数の範囲での因数分解は既知であり、左辺の \((x+\sqrt{-1})(x-\sqrt{-1})\) を適当に割ることで、\((a+b\sqrt{-1})(a-b\sqrt{-1})=p\) を得ることができます。

擬似コード

gauss_divisor = {}

def get_x(p):
  while True:
    y = random_int()
    x = modpow(y, (p-1)/4, p)
    if (x*x+1) % p == 0:
      return x

for p in primes:
  if p == 2:
    gauss_divisor[2] = 1+i
  else if p % 4 == 3:
    gauss_divisor[p] = p
  else:
    x = get_x(p)
    k = (x*x+1) / p
    g = x+i
    # g*conjugate(g) = k*p
    for q in prime_factor(k):
      qg = gauss_divisor[q]
      if (g*qg) % q == 0:
        g = g*qg / q  # g /= conjugate(qg)
      else:
        g = g*conjugate(qg) / q  # g /= qg
    gauss_divisor[p] = g