Project Euler 136

http://projecteuler.net/index.php?section=problems&id=136


前問と同じ方法では時間がかかりすぎます。前問では全体を同時に計算しましたが、これで遅いときは逆に個々に計算するとよいことがあります。ここでは解の個数が1になる条件を考えます。

n = (3d - a)(d + a)

なので、これがxyに次のように分解されるとしましょう。

x = 3d - a
y = d + a

すると、

d = (x + y) / 4
a = (3y - x) / 4

となるので、x + yは4の倍数、0 < x < 3yd, a自然数解を持つ条件になります。4の倍数が関係するので、n = 2e m(mは奇数)の形で考えます。
e = 0のとき、分解すると奇数同士になるので、4で割って余り1と3にならなければなりません。m = 1はダメです。m = p素数、以下同じ)なら、pが4の剰余が3なら(1, p)の一つのみ解となります。m = p2なら、(1, p2)と(p, p)が考えられますが、pの4の剰余が1でも3でも解になりません。m = q r(1 < q < r)なら、mの4の剰余が3でなければならないので、qrのどちらかが剰余1でどちらかが3です。少なくとも(q, r)、(1, q r)の2つが解になります。
e = 1のとき、奇数と偶数に分解されるのでダメ。
e = 2のとき、m = 1のとき(2, 2)が唯一の解。m = pなら、(2, 2p)が唯一の解。m = p2なら、(2p, 2p)と(2, 2p2)の2つが解、m = q r(1 < q < r)なら、少なくとも(2q, 2r)、(2, 2q r)の2つが解になります。
e = 3のとき、奇数と偶数に分解されたら解0個。偶数と偶数に分解されたら、4の剰余が0と2になるから解0個。
以下同様に議論を進めて、解が一つになるnの個数をまとめると、N未満のnを考えるとして、

e = 0のとき、N未満の4の剰余が3の素数の個数
e = 1のとき、0個
e = 2のとき、N/4未満の素数の個数
e = 3のとき、0個
e = 4のとき、N/16未満の素数の個数
eが5以上のとき、0個

となります。
下のように素直にコードを書いたら、1分を少し超えました。少し工夫したら切りましたが。

from itertools import *
from math import sqrt

def gen_primes(max_n):
    L = max_n / int(sqrt(max_n))
    M = max_n / L
    
    a = [ True ] * L
    for p in (n for n in xrange(2, L) if a[n]):
        for k in xrange(p * 2, L, p):
            a[k] = False
    primes = [ n for n in xrange(2, L) if a[n] ]
    for p in primes:
        yield p
    
    for m in xrange(1, M + 1):
        start = m * L
        end = start + L
        a = [ True ] * L
        for p in takewhile(lambda p: p * p < end, primes):
            k0 = (start + p - 1) / p * p - start
            for k in xrange(k0, L, p):
                a[k] = False
        
        for p in (start + k for k in xrange(L) if a[k]):
            if p <= max_n:
                yield p

N = 5 * 10 ** 7
print sum(1 for p in gen_primes(N - 1) if p % 4 == 3) \
    + sum(1 for p in gen_primes((N - 1) / 4)) \
    + sum(1 for p in gen_primes((N - 1) / 16))