Project Euler 21

Problem 21
(略)10000より小さいすべての友愛数の和を求めよ。
http://projecteuler.net/index.php?section=problems&id=21

10000より小さい自然数に対し、素因数分解をし約数の和を求める処理を2回繰り返せば友愛数かの判定ができます。


from itertools import ifilter

def make_primes(max_p): # 一番下のコード参照
def get_exp(n, p): # 同じく

def factorize(n):
f = [ ]
for p in primes:
if p * p > n:
break
e, n = get_exp(n, p)
if e:
f.append( (p, e) )
if n != 1:
f.append( (n, 1) )
return f

def d(n):
f = factorize(n)
return reduce(lambda x, y: x * y,
map(lambda t: (t[0] ** (t[1] + 1) - 1) / (t[0] - 1), f), 1) - n

def is_amicable(n):
m = d(n)
return n != d(n) and d(m) == n

t = time.clock()
N = 10 ** 5
primes = make_primes(int(N ** 0.5 + 0.5))
print sum(ifilter(is_amicable, xrange(2, N + 1)))
print time.clock() - t

しかし、次のように素因数分解の過程で直接約数の和を求めることもできます。以下は、12の約数の和を求めています。

s = 1 # 和
n = 12
p = 2 # 割り切れる素数
n = n / p2 = 3
s = s * (1 + p + p2) = 7
p = 3
n = n / 3 = 1
s = s * (1 + p) = 28

これで余分なコードが減って少し速くなります。


def sum_pow(n, p):
m = 1
while n % p == 0:
m = m * p + 1
n /= p
return m, n

def d(n):
n0 = n
s = 1 # sum of divisors
for p in primes:
if p * p > n:
break
m, n = sum_pow(n, p)
if m != 1:
s *= m
if n != 1:
s *= n + 1
return s - n0 # subtract self

しかし、ここまでのコードでは素因数分解が非常に遅くボトルネックになっています。高速化のためにエラトステネスのふるいと同様のアルゴリズムを使いましょう。例えば、10までで考えます。まず、10までの配列と、同じ長さの1を並べた配列を用意します。


1 2 3 4 5 6 7 8 9 10
1 1 1 1 1 1 1 1 1 1

2の倍数のとき、2が1個なら上を2で割り下を3倍、2が2個なら上を4で割り下を7倍、などとします。


1 1 3 1 5 3 7 1 9 5
1 3 1 7 1 3 1 15 1 3

3についても同様に処理します。


1 1 1 1 5 1 7 1 1 5
1 3 4 7 1 12 1 15 13 3

上の段で1でないものは素数なので、下にその素数+1を掛けます。


1 1 1 1 5 1 7 1 1 5
1 3 4 7 6 12 8 15 13 18

最後に自分自身を引きます。


1 1 1 1 5 1 7 1 1 5
1 1 1 3 1 6 1 7 4 8

この手順で約数の和をまとめて求めます。ただし、約数の和が求めた範囲より大きくなる場合があるので、その場合は上の単純な方法で計算します。これで、1桁以上速くなりました(恐らくC++ならもっと速くなります)。


from itertools import ifilter
import time

def make_primes(max_p):
if max_p <= 2:
return [ 2 ]
else:
a = [ True ] * (max_p + 1)
sqrt_max_p = int((max_p + 0.5) ** 0.5)
primes = make_primes(sqrt_max_p)
for p in primes:
for k in xrange(p, max_p + 1, p):
a[k] = False

for k in xrange(sqrt_max_p + 1, max_p + 1):
if a[k]:
primes.append(k)
return primes

def sum_divisors(max_n):
a = range(max_n + 1)
f = [ 1 ] * (max_n + 1)
for p in primes:
for k in xrange(p, max_n + 1, p):
m, a[k] = sum_pow(a[k], p)
f[k] *= m

for k in xrange(2, max_n + 1):
if a[k] != 1:
f[k] *= a[k] + 1
f[k] -= k

return f

def sum_pow(n, p):
m = 1
while n % p == 0:
m = m * p + 1
n /= p
return m, n

def calc_d(n):
n0 = n
s = 1 # sum of divisors
for p in primes:
if p * p > n:
break
m, n = sum_pow(n, p)
if m != 1:
s *= m
if n != 1:
s *= n + 1
return s - n0 # subtract self

def is_amicable(n):
m = d[n]
if m > M:
return calc_d(m) == n
else:
return m != n and d[m] == n

t = time.clock()
N = 10 ** 6
M = N
primes = make_primes(int(M ** 0.5 + 0.5))
d = sum_divisors(M)
print sum(ifilter(is_amicable, xrange(2, N + 1)))
print time.clock() - t