Python で多次元リストを作る

Python

Python で 2 次元リストを作るには、リスト内包表記を使って以下のように書くことが出来ます。

>>> rows, cols = 2, 3
>>> [[0] * cols for _ in range(rows)]
[[0, 0, 0], [0, 0, 0]]

2 次元まではいいのですが、3 次元以上になるとリスト内包表記で書くのはなかなか難しいです。
そこで、多次元リストを返す関数を作ってみました。

def multilist(ds, default):
    assert len(ds) > 0 and all(e > 0 for e in ds)

    def rec(index, parent):
        if index + 1 == len(ds):
            parent.extend(default() for _ in range(ds[-1]))
        else:
            for _ in range(ds[index]):
                child = []
                parent.append(child)
                rec(index + 1, child)

    root = []
    rec(0, root)
    return root

こういう風に使います。

>>> multilist([2, 3], lambda: 0)
[[0, 0, 0], [0, 0, 0]]
>>> multilist([2, 3, 4], lambda: 0)
[[[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]], [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]]

この関数を使って、ABC 029 D - 1 を解いてみました:

collections.defaultdict とタプルを使う方法もあります

>>> import collections
>>> d = collections.defaultdict(lambda: 0)
>>> d[1, 2, 3] = 100
>>> d[1, 2, 3]
100
>>> d[0, 0, 0]
0