Project Euler 135

プロジェクトオイラー
http://projecteuler.net/index.php

Q135.
自然数nについて、自然数x,y,zが等差数列になっていて、x2 - y2 - z2 = nを満たすとする。このような解が10ある100万より小さいnの個数。

d = y - zとして、

(z + 2d)2 - (z + d)2 - z2 = n
(3d - z)(d + z) = n

だから、nを分解すればよい。aとbに分解できるとすれば、

a = 3d - z
b = d + z
d = (a + b) / 4
z = (-a + 3b) / 4

より、

a + b ≡ 0(mod 4)
3b > a

が条件となる。よく見ると、nは4で割って余り2はダメ。また、素因数は重複を含めて4つ以上必要。
最初、nを1から増分して計算していったが、遅くて話にならなかったので、逆に素因数分解のほうからはじめた。100万より小さくなるような素因数分解を生成し、それについて条件を満たす約数の数を数えていった。



def gen_primes(p = 1):
if p == 1:
yield 2
if p <= 2:
yield 3
n = 3
if p > 4:
n = (p - 1) / 2 * 2 + 1
while True:
n += 2
p = 3
while True:
if p * p > n:
yield n
break
if n % p == 0:
break
p += 2

def gen_composites(n, q = 1):
yield [ ]
for p in gen_primes(q):
if q == 1:
if p ** 4 > n:
break
else:
if p > n:
break

i = 1
pow = p
while pow <= n:
for c in gen_composites(n / pow, p):
yield [ p, i ] + c
i += 1
pow *= p

def comp_num(c):
n = 1
i = 0
while i < len(c):
n *= c[i] ** c[i+1]
i += 2
return n

def num_div(c):
n = 1
for i in range(1, len(c), 2):
n *= c[i] + 1
return n

def gen_divisors(c):
n = comp_num(c)
num_primes = len(c) / 2
e = [ 0 ] * num_primes
d = 1
yield (1, n)
while True:
k = 0
while k < num_primes:
e[k] += 1
if e[k] <= c[k*2+1]:
d *= c[k*2]
break
e[k] = 0
d /= c[k*2] ** c[k*2+1]
k += 1
if k == num_primes:
break
yield (d, n / d)

def is_valid_div(a):
return (a[0] + a[1]) % 4 == 0 and a[1] * 3 - a[0] > 0

def C(c):
if len(c) == 0 or c[0] == 2 and c[1] == 1:
return 0
if num_div(c) < 10:
return -1

return sum(map(is_valid_div, gen_divisors(c)))

N = 10 ** 6
print sum(map(lambda c: C(c) == 10, gen_composites(N)))