Project Euler 171

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

Q171.
各桁の自乗和が自乗になる数で10²⁰より小さいものの総和の最後の9桁。

自乗和のほうから考える。20桁までなので、20*9*9以下の自乗数を考えればよい。それぞれについて、自乗による分割を求める。ただし、数字の組合せで考えて、各桁の数字が下位に行くにつれて大きくなるものだけ求める。求まったら、例えば、3桁で考えて122なら、最後の桁が2になる数は、122・212になる。これは、2は固定して残り2つのなかから1つ選ぶから、2!/1!1!となる。ある桁の和は、2*2!/1!1! + 1*2!/2! = 5となる。どの桁も同じだから、総和は3桁まで見ると、5の111倍ということになる。

from math import sqrt

import time
def g(n, m, l):

    if m > 0 and n <= m * m * l:

        for k in range(1, m + 1):

            r = n - k * k

            if r == 0:

                yield k

            elif r > 0:

                for e in g(r, k, l - 1):

                    yield k + e * 10
# 244 -> [ 17, 0, 1, 0, 2, ... ]

def encode(n):

    a = [ 0 ] * 10

    for i in range(N):

        a[n%10] += 1

        n /= 10

    return a
def sum2(n):

    s = 0

    a = encode(n)

    for i in range(1, 10):

        if a[i]:

            m = 1

            k = N - 1

            for j in range(10):

                rep = a[j]

                if i == j:

                    rep -= 1

                for l in range(rep):

                    m *= k

                    m /= l + 1

                    k -= 1

            s += m * i

    return s * (10 ** N - 1) / 9
def f_inverse(n):

    return sum(map(sum2, g(n, 9, N)))
t = time.clock()

N = 20

M = 9 * 9 * N

print sum(map(f_inverse, map(lambda x: x * x, range(1, int(sqrt(M + 0.5)) + 1))))

print time.clock() - t