Project Euler 211

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

Q211.
nの約数の平方の総和をσ(n)と書く。0 < n < 64000000 でσ(n)が平方数になるnの総和を求めよ。

nの素因数分解を

n = p₁^e₁...p_m^e_m

とすると、

σ(n) = (1 + p₁ + ... + p₁^e₁)...(1 + p_m + p_m^e_m)

となるから、最初、自然数を素因数分解するのではなく、直接素因数分解を生成するコードを書いていた。しかし、そのうちσ(n)（とn）を直接生成すればなんとか計算できることに気がついた。

from math import sqrt

from itertools import count
def is_square(n):

    m = int(sqrt(n + 0.5))

    return m * m == n
def is_prime(n):

    if n < 2:

        return False

    for p in primes:

        if p * p > n:

            return True

        elif n % p == 0:

            return False

    return True
def make_primes(n):

    m = int((n + 0.5) ** 0.5)

    for p in xrange(3, m + 1, 2):

        if is_prime(p):

            primes.append(p)

    

    a = [ 0 ] * ((n + 15) / 16)

    for p in primes:

        if p * p > n:

            break

        if p == 2:

            continue

        for k in xrange(3 * p, n, 2 * p):

            a[k>>4] |= 1 << (k & 15)

    

    for k in xrange((m + 1) / 2 * 2 + 1, n + 1, 2):

        if (a[k>>4] & (1 << (k & 15))) == 0:

            primes.append(k)
def gen_sigma(n, i = 0):

    yield (1, 1)

    

    for k in xrange(i, len(primes)):

        p = primes[k]

        if p > n:

            break

        m = n

        sigma = 1

        for j in count(1):

            m /= p

            if m == 0:

                break

            sigma = sigma * p ** 2 + 1

            for e in gen_sigma(m, k + 1):

                yield (p ** j * e[0], sigma * e[1])
N = 64000000

primes = [ 2 ]

make_primes(N)
s = 0

for n, sigma in gen_sigma(N):

    if is_square(sigma):

        s += n

print s