Project Euler 135

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

初項をa、公差をdとすると、

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

となります。
dを固定してaを変えましょう。dが小さいとき、1 ≤ a ≤ 3d - 1となりますが、dが大きいときにa = dをピークとしてnが100万以上になるので、そうでない領域が分断されます。こういうときのためにiterableをつなぐchainという関数がPythonには用意されています。これで一つのiterableになるので、dが小さいときと同じ扱いができます。

from itertools import *

def f(d, a): return (3 * d - a) * (d + a)

N = 10 ** 6
L = 10
b = [ 0 ] * N
for d in xrange(1, (N + 1) / 4 + 1):
    if 4 * d * d < N:
        g = (f(d, a) for a in xrange(1, 3 * d))
    else:
        g1 = takewhile(lambda n: n < N, (f(d, a) for a in count(1)))
        g2 = takewhile(lambda n: n < N, (f(d, a) for a in
                                        xrange(3 * d - 1, 0, -1)))
        g = chain(g1, g2)
    for n in g:
        b[n] += 1

print sum(1 for n in xrange(N) if b[n] == L)