第２問「奇数和平方数」問題解答

「奇数和平方数」問題の解法です。
問題文はこちら。
他の問題の解法は、第１問、第３問、第４問、第５問。

連続する奇数の和が平方数になるものを探す問題です。
1 から m 番目の奇数の和については、有名な公式が存在します。イコール m² です。
ただ今回は、1 以外の奇数から始まるものについても考慮しなければなりません。
始点の一つ前の奇数を p 番目の奇数、終点を q 番目の奇数とするならば、この間にある奇数の和は q² - p² となります。
これが平方数になるということは・・要するに、ある自然数 r が存在して、

$\begin{eqnarray} p^2 - q^2 = r^2 \\ \Leftrightarrow q^2 + p^2 = p^2 \end{eqnarray}$

となる (q, r, p) の組、すなわちピタゴラス数を求める問題に帰着されます。

ピタゴラス数の列挙といえば定番ネタです。
ピタゴラスの定理 - Wikipedia の記事に、原始ピタゴラス数（q, r, p が互いに素であるピタゴラス数）を列挙するやり方が載ってます。互いに素な自然数の組 (s, t) で、かつ s＞t、s－t が奇数であるものについて、

$\begin{eqnarray} (s^2-t^2, 2st, s^2+t^2) または (2st, s^2-t^2, s^2+t^2) \end{eqnarray}$

は必ず原始ピタゴラス数になります。これでまずは原始ピタゴラス数を列挙しましょう。

題意を満たすのは原始ピタゴラス数だけではありません。実際には、これの定数倍もすべてが題意を満たします。そうしたものの個数は、p 番目の奇数 2p-1 が n 以下であることに注意して、(n+1)/p/2 であると分かります。

コード例は上の通りです。q と r の入れ替わりを考慮してループ内で足す数を 2 倍している点と、最後に p = 0 の場合を足しこんでいる点に注意して下さい。

def gcd(x, y):
	return gcd(y % x, x) if x > 0 else y

def f(n):
    answer = 0
    for s in range(1, n):
        if s*s > n: break
        for t in range(1, s):
            p = s*s + t*t
            if p > n: break
            if (s - t) % 2 == 0: continue
            if gcd(s, t) > 1: continue
            answer += 2*((n+1)/p/2)
    answer += (n+1)/2
    return answer

print f(10000000)

このコードで ideone で 3 秒程度です。

ちなみに原始ピタゴラス数を列挙するには、魔法の行列を使う、もっと高速な方法があります。
ピタゴラスの数を生む行列を参照して下さい。

こちらは C++ で書きました。C++ の速さも相まって、ideone で 0.1 秒切りますね。

#include <iostream>
using namespace std;
typedef long long int	s64;

s64 Do(s64 q, s64 r, s64 p, s64 n) {
	s64 a = 2*((n+1)/p/2);
	if (a == 0) return 0;
	a += Do(q-2*r+2*p, 2*q-r+2*p, 2*q-2*r+3*p, n);
	a += Do(q+2*r+2*p, 2*q+r+2*p, 2*q+2*r+3*p, n);
	a += Do(-q+2*r+2*p, -2*q+r+2*p, -2*q+2*r+3*p, n);
	return a;
}

int main() {
	s64 n = 10000000;
	s64 r = Do(3, 4, 5, n) + (n+1)/2;
	cout << r << endl;
}