Project Euler 45（1）

Problem 45
三角数、五角数、六角数は次の公式で生成される。
T_n = n(n+1)/2 1, 3, 6, 10, 15, ...
P_n = n(3n-1)/2 1, 5, 12, 22, 35, ...
H_n = n(2n-1) 1, 6, 15, 28, 45, ...
T₂₈₅ = P₁₆₅ = H₁₄₃ = 40755 であることが確かめられる。
この次の五角数でも六角数でもある三角数を見つけよ。
http://projecteuler.net/index.php?section=problems&id=45

六角数は三角数でもあるので、三角数は考える必要はありません。
五角数と六角数を同時に生成して、五角数のほうが小さければ五角数を生成、六角数のほうが小さければ六角数を生成、等しければ出力して両方生成します。これは簡単でも有用なアルゴリズムなので覚えましょう。
しかし、この次を出そうとすると時間がかかります。

from itertools import count, islice
def gen_polygonal(p):

    for n in count(1):

        yield n * ( (p - 2) * n - p + 4) / 2
def gen_both():

    gp = gen_polygonal(5)

    gh = gen_polygonal(6)

    p = gp.next()

    h = gh.next()

    while True:

        if p == h:

            yield p

            p = gp.next()

            h = gh.next()

        elif p < h:

            p = gp.next()

        else:

            h = gh.next()
print sum(islice(gen_both(), 2, 3))