Project Euler 181

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

Q181.
3つのBと1つのWの分割を考えると、次のような7通りが考えられる。
(BBBW), (B,BBW), (B,B,BW), (B,B,B,W), (B,BB,W), (BBB,W), (BB,BW)
60のBと40のWなら何通りか。

これは二次元版の分割数というべきものである。一次元の分割数と同様の母関数を作る。最初の例なら、

Q(p, q) = (1 + p + p² + p³)(1 + p²)(1 + p³)(1 + q)(1 + pq)(1 + p²q)(1 + p³q)

このp³qの係数を拾えばよい。
最初素直に組んだらPythonでの多項式の積の計算が遅くて1時間以上かかりそうだったが、工夫して40秒にした。

from itertools import product
def mul(a, k1, k2):

    c = [ [ 0 ] * (M + 1) for i in range(N + 1) ]

    for n1 in xrange(N + 1):

        for m1 in xrange(M + 1):

            for n2, m2 in gen(N, M, k1, k2):

                n = n1 + n2

                m = m1 + m2

                if n <= N and m <= M:

                    c[n][m] += a[n1][m1]

    return c
def gen(n, m, k1, k2):

    if k1 == 0:

        m2 = 0

        while m2 <= m:

            yield (0, m2)

            m2 += k2

    elif k2 == 0:

        n2 = 0

        while n2 <= n:

            yield (n2, 0)

            n2 += k1

    else:

        n2 = 0

        m2 = 0

        for i in xrange(min(n / k1, m / k2) + 1):

            yield (n2, m2)

            n2 += k1

            m2 += k2
def make_poly(n, m, k1, k2):

    if k1 == 0 and k2 == 0:

        end = 1

    elif k1 == 0:

        end = m / k2 + 1

    elif k2 == 0:

        end = n / k1 + 1

    else:

        end = min(n / k1, m / k2) + 1

    

    for i in range(end):

        a[i*k1][i*k2] = 1

    return a
N = 60

M = 40

a = [ [ 0 ] * (M + 1) for i in range(N + 1) ]

a[0][0] = 1

for k1, k2 in product(range(N + 1), range(M + 1)):

    if k1 == 0 and k2 == 0:

        continue

    elif k1 == 0 and k2 == 1:

        a = make_poly(N, M, k1, k2)

    else:

        a = mul(a, k1, k2)

print a[N][M]