2015年11月10日

itertools.productで2重ループ

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


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

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

この行列の各要素を全走査するにはどのようにすればよいでしょう?
簡単に思い浮かぶのは、以下のような2重ループだと思います。
for i in range(3):  # 行のループ
    for j in range(3):  # 列のループ
        print(matrix[i][j])

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


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

from itertools import product

for i, j in product(range(3), repeat=2):
    print(matrix[i][j])
見事にインデントが1つで済みました。Coolですね!



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

from itertools import product

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

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

0 件のコメント:

コメントを投稿