Project Euler 136

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

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

注意深く調べると、pを素数として、nは4の剰余が3のp、4p、16p、4、16となる。
時間がかかるのは素数の列挙なので、それを少し工夫した。

def make_prime(n):

    m = 3

    while m <= n:

        if is_prime(m):

            primes.append(m)

        m += 2
def is_prime(n):

    for p in primes:

        if p * p > n:

            return True

        if n % p == 0:

            return False

    return True
def gen_primes(N):

    for p in primes:

        yield p

    

    n = p + 2

    while n < N:

        if is_prime(n):

            yield n

        n += 2
N = 5 * 10 ** 7

primes = [ 2 ]

make_prime(int(N ** 0.5 + 0.5))

counter = 2

for p in gen_primes(N):

    if p == 2:

        continue

    if p % 4 == 3:

        counter += 1

    if p * 4 < N:

        counter += 1

        if p * 16 < N:

            counter += 1

print counter