[Python] product関数を使って多重ループを1つのループにする

多重ループを書くとネストが深くなり、読みにくさが倍増する。

関数を使うなど、ループを浅くする方法はいくつかあるが、ここではPython標準ライブラリのitertoolsライブラリを使う方法を紹介する。

例えば次のようなコードがあったとしよう。

コード 1:

for i in range(3):
    for j in range(2):
        print([i, j])

出力 1:

[0, 0]
[0, 1]
[1, 0]
[1, 1]
[2, 0]
[2, 1]

このコードには二重ループが含まれている。二重ループであればそれほど読みにくさは感じないが、他の処理と組み合わさって三重以上のループになってしまった場合、読みにくさを感じる人も多いのではないか。

そこで、Python標準ライブラリであるitertoolsライブラリに用意されている、product関数を使って、この二重ループを一重ループにする。
docs.python.org

コード 2:

from itertools import product

for i, j in product(range(3), range(2)):
    print([i, j])

出力 2 (出力1と同じ):

[0, 0]
[0, 1]
[1, 0]
[1, 1]
[2, 0]
[2, 1]

product関数は、複数のリストまたはジェネレータを受け取り、それらのすべての組み合わせを生成するジェネレータを返す。これにより、シンプルな二重ループを1つのループに落とすことができる。

同様にして、三重ループも1つのループに落とすことができる。

コード 3（出力結果は省略）:

from itertools import product

for i, j in product(range(3), range(2), range(4)):
    print([i, j, k])

product関数のrepeat引数に整数を指定すれば、与えたリストまたジェネレータを何個分組み合わせるかを指定できる。下の例では、0から2までを生成するジェネレータを二つ組み合わせて、合計9通りの組み合わせを生成するジェネレータを返している。

コード 4:

from itertools import product

for i, j in product(range(3), repeat=2):
    print([i, j])

出力 4:

[0, 0]
[0, 1]
[0, 2]
[1, 0]
[1, 1]
[1, 2]
[2, 0]
[2, 1]
[2, 2]

itertoolsライブラリにはproduct関数の他にも有用な関数が用意されているので、興味を持った方は一度ドキュメントを覗いてみるのがよいと思う。
docs.python.org