itertools.productで２重ループ

２重ループを書かなくて済む便利なツール itertools.product について紹介します。

以下のような２重配列による正方行列があったとします。

matrix = [
    [6, 1, 8],
    [7, 5, 3],
    [2, 9, 4],
]

この行列の各要素を全走査するにはどのようにすればよいでしょう？
簡単に思い浮かぶのは、以下のような２重ループだと思います。

for i in range(3):  # 行のループ
    for j in range(3):  # 列のループ
        print(matrix[i][j])

上の書き方で、動作自体は全く正しい物となります。
しかし、これだけ単純なループのために、インデントを２つも下げるのはCoolとは言えません。Python ならば２重ループを１行で実現する手段があるのでは、と考えてします。

ビルドインにはありませんが、itertools.product というモジュールを使うことで、２重ループを１行で実現できます。

from itertools import product

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

見事にインデントが１つで済みました。Coolですね！

itertools.product は様々な多重ループに対応できるようになっています
以下に２つの例を挙げます。

from itertools import product

# ３重ループ
for i, j, k in product(range(3), repeat=3):
    print(i, j, k)

# １重目と２重目で異なるループ
for i, j in product(range(12), range(31)):
    print(i, j)