for ループが多く現れ始めたら itertools の活用を考える

Python の for ループとネスト

Python でコードを書いている時に、for ループ内でさらに for ループを回したくなるときってありますよね。

たとえば、

# 2023年から2024年までの年月の組みあわせを全部出す
for year in range(2023, 2025):
    for month in range(1, 13):
        print(year, month)

とすれば

2023 1
2023 2
(中略)
2024 11
2024 12

とまあ、こんな感じです。

こういった書き方の課題はやはり、Pythonのネスト(インデント)が深まっていって可読性が下がっていく可能性があることですね。

こんなときに使えるのが itertools です。

ぜひ覚えておきたい標準ライブラリの一つだと思いますので紹介します。

itertools とは

公式ドキュメント「itertools --- 効率的なループ用のイテレータ生成関数群」の通り、

ループで効率的にイテレータ(公式ドキュメントでいうところの「データの流れ」)を作るための標準ライブラリです。

先ほどの二重ループは itertools.product を使うことによって

import itertools
for year, month in itertools.product(range(2023, 2025), range(1, 13)):
    print(year, month)

こう書けます。

個人的には、 for ループ あるいは その元となるイテレータ が妙に複雑化しはじめたりするとまず思いを馳せる標準ライブラリ です。

公式ドキュメントには

プログラム言語 APL, Haskell, SML からアイデアを得ていますが、 Python に適した形に修正されています。

と書いてあって、そうだったのか…という感じです。(残念ながらいままでHaskellに触れることはなかったので…)

よく使うitertools

ここからは、よく使うitertoolsをまとめておきます。

詳しく知りたい方は先ほどの公式ドキュメント「itertools --- 効率的なループ用のイテレータ生成関数群」を参照されると良いと思います。

itertools.product

まずダントツで使います。先ほどのような多重ループが発生すると、まずこれで解決です。

あと、 repeat という引数があります。これも地味に便利なのでぜひ覚えておきたいところです。

同じイテレータを repeat に指定した回数食わせた場合と同じ結果を返してくれます。

import itertools
for x, y, z in itertools.product(range(2), repeat=3):
    print(x, y, z)

のアウトプットは

0 0 0
0 0 1
0 1 0
0 1 1
1 0 0
1 0 1
1 1 0
1 1 1

となります。

itertools.permutations

いわゆるパーミュテーションです。

リストの中から指定した数を重複なく抜き出してきてくれます。

list(itertools.permutations(range(3), 2))

の出力は

[(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)]

ということです。数学的に言うと 3P2 = 3 * 2 = 6 個出てくる、みたいなことですね。

itertools.combinations

こちらは組みあわせですので、パーミュテーションから順不同としたときの重複を除いたものです。

list(itertools.combinations(range(3), 2))

は

[(0, 1), (0, 2), (1, 2)]

となります。 3C2=3 個でてきてます。

同じ数字による (1, 1) のような組みあわせを認める combinations_with_replacement というのもあります。

>>> list(itertools.combinations_with_replacement(range(3), 2))
[(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)]

ちなみにこれはほぼ使ったことはありません…

ここぞというときに使えると気持ちよさそうです。

基本はこの3つで大丈夫だと思います。

これらは組みあわせ系の関数ですが、実際は accumulate() のように自分で思わず実装してしまいそうな処理も含まれていたりもしますので、一度目を通しておくことをオススメします。

>>> list(itertools.accumulate([1, 2, 3, 4, 5]))
[1, 3, 6, 10, 15]

ただし、中にはリスト内包表記でもいいのでは？となる例もあります。

repeat() の例に挙がっている

>>> list(map(pow, range(10), itertools.repeat(2)))
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

などは、range(10) の0〜9に対して、 pow(0, 2), pow(1, 2), … pow(9, 2) したものを返す、ということで、格好いいのですが…

リスト内包表記で

>>> [pow(n, 2) for n in range(10)]
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

とも書けるものであり、こちらの方が読む側の見通しは良いようにも思います。

一方で、前者の書き方のほうが map で逐次的に処理を行うため、今回のような range(10) などではなく、もっと大きなイテレータを相手に処理を行う場合には使える書き方とも言えるでしょう。

つまり状況に応じた使い分けをしたほうが良さそうなものもある ということですね。何事にも言えることですが…

まとめ

今回はPythonの中でも使えるシーンは割と多い標準ライブラリ itertools について扱いました。

特にネストが深くなってしまっている for ループなどが目の前に現れたら思い出して使ってみていただければと思います。

毎度のことながら、標準ライブラリの世界は奥が深いですね…掘れば掘るほどお宝が出てくる感じがあります。