PythonのライブラリをC++で作成する（1）

PythonのライブラリをC++で作ります。
例えば、次のような問題を考えます。

H×Wの長方形があり、そこに1×2の長方形を隙間なく敷き詰めます。敷き詰め方は何通りあるでしょう。

2×3なら、

の2通りです。
さて、この問題はどのように解くでしょう。この手の問題は、状態遷移を考えます。左から順に詰めていきます。そして、縦に見て最も左が空いている場所に長方形を埋めます。最も左が空いている場所が複数あれば、上から詰めます。空の状態から一つ長方形を詰めた状態は次の2つのみです。

この時、全て埋まっている列は消した状態を考えます。H = 3で考えて、例えば、

のようにします。これを(2, 0, 1)、もしくはもう少し簡単に201と表します。000から始まって、HW/2個長方形を敷いて000になっていたらH×Wが隙間なく敷き詰められたことになります。

状態遷移は、

000 -> 200, 110
200 -> 220, 100
220 -> 000
110 -> 001
100 -> 120, 000
001 -> 201, 000
120 -> 011
201 -> 110
011 -> 100

となり、これを行列にして、この行列のW/2乗を計算すればいいわけです。

コードは次のようになりました。

# coding: utf-8
from collections import Counter
import sys
import time


#################### Matrix ####################

def mat_mul(A, B, D):
    H, M, W = len(A), len(B), len(B[0])
    C = [ [0] * W for _ in range(H) ]
    for i in range(H):
        for j in range(W):
            C[i][j] = int(sum(A[i][k] * B[k][j] for k in range(M)) % D)
    return C

def mat_pow(M, e, D):
    if e == 1:
        return M
    elif e % 2 == 1:
        return mat_mul(M, mat_pow(M, e-1, D), D)
    else:
        A = mat_pow(M, e/2, D)
        return mat_mul(A, A, D)


#################### process ####################

def f(H, W, D):
    def normalize(diagram):
        min_l = min(diagram)
        return tuple(row - min_l for row in diagram)
    
    def nexts(diagram):
        y = next(y for y, row in enumerate(diagram) if row == 0)
        yield diagram[:y] + (2,) + diagram[y+1:]
        if y+1 < H and diagram[y+1] == 0:
            yield diagram[:y] + (1, 1) + diagram[y+2:]
    
    def walk(diagram):
        for diagram1 in nexts(diagram):
            diagram2 = normalize(diagram1)
            if diagram2 not in set_diagram:
                set_diagram.add(diagram2)
                walk(diagram2)
    
    def make_graph(H):
        graph = Counter()       # { (orig, dest): counter }
        for diagram in set_diagram:
            for diagram1 in nexts(diagram):
                graph[(diagram, normalize(diagram1))] += 1
        return graph
    
    def make_matrix(H):
        graph = make_graph(H)
        M = [ [ graph[(orig, dest)] for orig in set_diagram ]
                                    for dest in set_diagram ]
        row_to_index = dict((diagram, i) for i, diagram in enumerate(set_diagram))
        return (M, row_to_index)
    
    init_diagram = (0,)*H
    set_diagram = set([init_diagram])
    walk(init_diagram)
    print(len(set_diagram))
    M, row_to_index = make_matrix(H)
    M_pow = mat_pow(M, W*H/2, D)
    i = row_to_index[init_diagram]
    return M_pow[i][i]


#################### main ####################

t0 = time.time()
D = 10**9+7
H = int(sys.argv[1])
W = int(sys.argv[2])
print(f(H, W, D))
print(time.time() - t0)

これを9×100, D = 10**9 + 7で計算してみましょう。

1134
256330968
1756.864093542099

最初の1134は、行列の大きさです。約30分かかっていますね。競技プログラミングでは、こんなに行列が大きくなることはふつうないです。
さて、計算時間はこの行列の掛け算がほとんど全てです。この手のリストを使う計算は、PyPyにすると速くなります。

1134
256330968
277.08021235466003

6.3倍速くなりました。
Project Eulerでは、PyPyは事実上の標準です。PyPyを使いましょう。
もっと速くするために、C++によるライブラリを作ります。