Project Euler 217

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

Q217.
正の整数が、10進で表して上位半分と下位半分の各桁の和が等しいとき、平衡数と呼ぶ。10⁴⁷より小さい平衡数はいくつあるか？3¹⁵の剰余で答えよ。

これはやさしい。
半分の桁の和がある数を取る和と個数を計算していけばいい。例えば、半分が2桁で各桁の和が2なら、20と11を取りうるから、U1(2, 2) = (31, 2)とする。今のは上位半分の場合で、下位半分の場合なら02も加わるので、U2(2, 2) = (33, 3)となる。これらは漸化式で簡単に求まる。上位半分と下位半分は、漸化式は同じで、初期値がわずかに違うだけである。これらが求まったら、あとは適当に数える。

def calc_next_term(P):

    n = (len(P) - 1) / 9 + 1

    U = [ 0 ] * (n * 9 + 1)

    for m in range(n * 9 + 1):

        a = range(max(0, m - n * 9 + 9), min(9, m) + 1)

        u = sum(map(lambda k: 10 * P[m-k][0] + k * P[m-k][1], a))

        v = sum(map(lambda k: P[m-k][1], a))

        U[m] = (u, v)

    return U
N = 47

S = [ 0 ] * (N + 1)

U1 = [ (m, 1) for m in range(10) ]

U2 = [ (m, 1) for m in range(10) ]

U1[0] = (0, 0)
for n in range(1, N + 1):

    if n == 1:

        S[n] = 45

    else:

        m = n / 2

        a = range(m * 9 + 1)

        if n % 2 == 0:

            if n > 2:

                U1 = calc_next_term(U1)

                U2 = calc_next_term(U2)

            S[n] = sum(map(lambda k: U1[k][0] * U2[k][1], a))

            S[n] *= 10 ** m

            S[n] += sum(map(lambda k: U2[k][0] * U1[k][1], a))

        else:

            S[n] = sum(map(lambda k: U1[k][0] * U2[k][1], a))

            S[n] *= 10 ** (m + 1)

            S[n] += sum(map(lambda k: U2[k][0] * U1[k][1], a))

            s = sum(map(lambda t: t[0][1] * t[1][1], zip(U1, U2)))

            S[n] = S[n] * 10 + 45 * 10 ** m * s

print sum(S) % 3 ** 15