Project Euler 44（1）

Problem 44
五角数は次の公式で生成される：P_n = n(3n - 1) / 2。
（中略）五角数のペアP_j, P_kで、その差と和も五角数となり、D = |P_k - P_j|が最小となるものを見つけよ。そのDの値は？
http://projecteuler.net/index.php?section=problems&id=44

まともに考えると、k > j として、P_k - D が五角数なら、あとは P_j + P_k が五角数であることをチェックすればよいです。
同じDに対してどこまでのkについてチェックすればよいでしょう。D = P_k - P_k-1のときが最大のkになります。

P_k - P_k-1 = 3k - 2

となるので、

k ≤ (D + 2) / 3

となります。
nが五角数になるかの判定は、

P_m = m(3m - 1) / 2 = n
3m² -m - 2n = 0
m = (1 + √(1 + 24n)) / 6

だから、1 + 24n が平方数であることをチェックして、その平方根が6で割って5余るかを調べることになります。

from itertools import imap, count

from math import sqrt
def is_pentagonal(n):

    p = 1 + 24 * n

    q = int(sqrt(p))

    return q * q == p and q % 6 == 5
def P(n):

    return n * (3 * n - 1) / 2
def gen_pentagonal_pair():

    for i in count(1):

        Pd = P(i)

        for k in xrange(i + 1, (Pd + 2) / 3):

            Pk = P(k)

            yield Pd, Pk
def solve():

    for Pd, Pk in gen_pentagonal_pair():

        if is_pentagonal(Pk - Pd) and is_pentagonal(Pk * 2 - Pd):

            return Pd
print solve()

しかし、これはいつまで経っても返ってきません。計算量が膨大なためです。
どうすればよいでしょう。例えば、D = P₁₀ = 145 は k ≤ 49まで調べなければなりません。最大のkは P_k - P_k-1 = P₁₀ から出ましたが、P_k - P_k-2 = P₁₀とすると、

P_k - P_k-1 = 6k - 7 = 145
k = 152 / 6 = 30.…

となります。すなわち、31から48は調べなくてもよかったわけです。

P_k - P_k-q = P_d

を調べることにしましょう。q < d です。q = dだと、k = d となるからです。

P_k - P_k-q = 3qk - q(3q + 1) / 2 = P_d

を満たす自然数kが存在すれば、P_d, P_j, P_kが求まり、あとは P_j + P_k が五角数かをチェックします。

from itertools import imap, count

from math import sqrt
def is_pentagonal(n):

    p = 1 + 24 * n

    q = int(sqrt(p))

    return q * q == p and q % 6 == 5
def P(n):

    return n * (3 * n - 1) / 2
def gen_pentagonal_triplet():

    for d in count(1):

        Pd = P(d)

        for q in xrange(1, d):

            if (Pd + q * (3 * q + 1) / 2) % (3 * q) == 0:

                k = (Pd + q * (3 * q + 1) / 2) / (3 * q)

                yield Pd, P(k - q), P(k)
def solve():

    for Pd, Pj, Pk in gen_pentagonal_triplet():

        if is_pentagonal(Pj + Pk):

            return Pd
print solve()