二平方の和の個数

今回はヤコビの二平方定理 (Wiki)という面白い定理を見つけたので、これを使うプログラムを書いてみました。以下のような問題を考えてみます。

問題：　$N=10^{12}$、$a^2+b^2=c^f (a,b,cは自然数,0<a<b, f\geq 3,\space c^f<N)$となる(a,b,c)の数を数えよ

これを素直にインプリすると、以下のようになりますが、Python(Google Colab)では約45分かかります。

from time import time

def count2sq(f,N):
  ret = 0
  for c in itertools.count(2):
    c_f = c**f
    if c_f > N: break
    for b in itertools.count(2):
      b2 = b**2
      if (1+b2) > c_f: break
      a2 = (c_f - b2)
      if a2 >= b2: continue
      a = int(a2**(1/2)+0.00001)
      if (a**2 == a2):
        ret += 1
  return ret

N, count = 10**12, 0
start = time()
for f in itertools.count(3):
  if 2**f > N: break
  count += count2sq(f,N)
print(f"**** count = {count},  {time() - start:.2f}sec ****")
# **** count = 15992,  2697.05sec ****

ヤコビの二平方定理

日本語版のヤコビの二平方定理 (Wiki)の説明だと分かりづらいので英語版のSum of squares functionの$k=2$の説明を引用します。

\begin{align}
&a^2+b^2=nを満たす整数のペア(a,b)の個数をr_2(n)を求める \\ \\
&nの素因数分解が以下のように表されるとすると、\\
&\space n=2^gp^{f_1}_1p^{f_2}_2\cdots q^{h_1}_1q^{h_2}_2\cdots\\
&\ \ \ \ ここで \\
&\ \ \ \ p_i\equiv 1\space \pmod 4\\
&\ \ \ \ q_i\equiv 3\space \pmod 4 \\
&r_2(n)は以下の式で求められます\\
\end{align}

\begin{eqnarray}
r_2(n)=\left\{ \begin{array}{ll}
4(f_1+1)(f_2+1)\cdots &(h_1,h_2,\cdotsがすべて偶数の時）\\
0 & (h_iの一つでも奇数の時)
\end{array} \right.
\end{eqnarray}

要は$n$の素因数を求めてmod 4が１になるものの数を数えればいいということになります。元の問題は$n=c^f$なので、$c$の素因数分解をして各々$f$倍すれば良いということでヤコビの二平方定理が効率よく利用できます。
ただしこの$r_2(n)$は$a,b$の各々の$\pm$と$a,b$の入れ替えを含んでいるので問題の数を求めるには$4\times 2=8$で割る必要があります。

Pythonでインプリメントすると以下のようになり、無事、同じ答えが0.10secで求まりました。

from sympy import factorint
import itertools
from time import time

def sumE2f(N, f):
  sum = 0
  for c in itertools.count(2):
    c_f = c**f 
    if c_f > N: break
    cfct = factorint(c)
    rn, p3 = 4, 0
    for p in cfct.keys():
      cfp = cfct[p]*f
      if p%4==3 and cfp % 2 == 1: rn = 0
      if p%4==1: rn *= (cfp+1)
    if rn > 4:
      sum += rn // 8
  return sum

start = time()
np, sum2 = 12, 0
N = 10**np
for f in itertools.count(3):
  if 2**f > N: break
  sm = sumE2f(N,f)
  if sm > 0:
  sum2 += sm 

print(f"**** count = {sum2},  {time() - start:.2f}sec ****")
#**** count = 15992,  0.10sec ****

このアルゴリズムはEuler Project: Problem 678等を解くのに役に立ちます。