Python​/itertoolsでループを効率化

itertoolsとは

プログラムを書いていると、よく多重ループを使う場面が出てきます。 例えば、3重ループなら

for i in range(2):
    for j in range(2):
        for k in range(2):
            print i,j,k

というような物です。 3重ぐらいなら特に手で書いても問題ありませんが、それ以上だとネスト（階層）が深くなって管理が大変です。または何重かわかっていない場合などはそもそも手で書く事はできません。

他の手続き型言語ではそのような場合、再帰関数を使うのが一般的で、Pythonでも同様のことが出来ます。 例えば以下のような関数で、上の多重ループと全く同じ処理が出来ます。

def loopFunction(maxLen,maxLoop,loopList = []):
    if len(loopList) < maxLoop:
        for i in range(maxLen):
            l = list(loopList)
            l.append(i)
            loopFunction(maxLen,maxLoop , l)
    else:
        for i in loopList:
            print i,
        print

loopFunction(2, 3)

この方法だと、好きな数の多重ループも出来ます。 しかし、見ての通り、ちょっとコードが複雑になりすぎてしまいます。

こんな時に、itertoolsを使うと、以下のようなシンプルなコードで同じ処理が可能です。

import itertools

for e in itertools.product(range(2),repeat=3):
    for i in e:
        print i,
    print

しかも「repeat=3」の部分を変更すれば、何重ループにも対応できます。

itertoolsは他にもいろいろな関数があり、さまざまな変数の組み合わせなどを作成してくれます。

このようにitertoolsは変数の生成によりループを効率化するモジュールです。

全ての組み合わせを生成する「product」

itertoolsの中でも、特によく使うのはproductでしょう。これは日本語で「積」を表す名前の通り、ベクトルやテンソルの解析で言うところの直積にあたるもので、引数に与えられた配列などの全ての可能な組み合わせを返します。

for i in itertools.product("AB","DE"):
    print i

結果:

('A', 'D')
('A', 'E')
('B', 'D')
('B', 'E')

つまり、以下のような通常のforの多重ループと同じです。

for i in "AB":
    for j in "DE":
        print i,j

単純なforの多重ループはたいていproductに置き換えることができます。そのため使う機会が多く、覚えておくとコードがかなりシンプルにできて便利です。

また、同じ配列を組み合わせる場合は、以下のようにrepeat変数を指定します。

for i in itertools.product("AB",repeat=3):
    print i

結果:

('A', 'A', 'A')
('A', 'A', 'B')
('A', 'B', 'A')
('A', 'B', 'B')
('B', 'A', 'A')
('B', 'A', 'B')
('B', 'B', 'A')
('B', 'B', 'B')

順列「permutations」

permutationsは要素の順列を返します。

for i in itertools.permutations("ABC"):
    print i

結果:

('A', 'B', 'C')
('A', 'C', 'B')
('B', 'A', 'C')
('B', 'C', 'A')
('C', 'A', 'B')
('C', 'B', 'A')

第二引数に整数を指定すると、その数だけ要素を取り出した並べ替えを返します。

for i in itertools.permutations("ABC",2):
    print i

結果:

('A', 'B')
('A', 'C')
('B', 'A')
('B', 'C')
('C', 'A')
('C', 'B')

順番に意味があり、('A', 'B')と('B', 'A')は別の物になります。 例えば、複数地点を経由した距離を、全ての場合で出したいというときに使えます。

組み合わせ「combinations」

combinationsは組み合わせを作成します。例えば「ABC」から2つ取り出す組み合わせは以下のようになります。

for i in itertools.combinations("ABC",2):
   print i

結果：

('A', 'B')
('A', 'C')
('B', 'C')

順列とは異なり、並べ替えて同じになる物は省略されます。つまり('A', 'B')と('B', 'A')は同じものとされてどちらかのみが返されます。

重複組み合わせ「combinations_with_replacement」

combinations_with_replacementは重複を許した組み合わせです。つまり同じ要素を何度取り出してもかまわない組み合わせです。

for i in itertools.combinations_with_replacement("ABC",2):
    print i

結果：

('A', 'A')
('A', 'B')
('A', 'C')
('B', 'B')
('B', 'C')
('C', 'C')

例えばいくつかの商品から2つだけ買っても良い場合、その値段を全て求めたいときなどに使えます。重複を許さないなら通常の組み合わせとなります。

まとめ

これ以外にもitertoolsには多くの機能があります。ただ、上のものに比べると、使用頻度が低かったり、他の記法で簡単に同じ事ができたりすることが多いです。とりあえず覚えるならこれらの組み合わせから始めるがおすすめです。