Python で for 文を書きがちな人！
itertools を知ってると、同じことがもっと短く、きれいに書けるかもしれません。
こんなコード書いてませんか...？
たとえばリストの中身を全部組み合わせたいとき、こうやって書いてませんか？
colors = ['red', 'blue']
sizes = ['S', 'M', 'L']
for color in colors:
for size in sizes:
print(f"{color} / {size}")
出力
red / S
red / M
red / L
blue / S
blue / M
blue / L
itertools を使うとこうなります！
from itertools import product
colors = ['red', 'blue']
sizes = ['S', 'M', 'L']
for color, size in product(colors, sizes):
print(f"{color} / {size}")
他にも...
data = ['A', 'B', 'C', 'D']
use_flags = [1, 0, 1, 0]
result = []
for d, flag in zip(data, use_flags):
if flag:
result.append(d)
print(result)
出力
['A', 'C']
itertools を使うとこうなります！
from itertools import compress
data = ['A', 'B', 'C', 'D']
use_flags = [1, 0, 1, 0]
print(list(compress(data, use_flags)))
itertools の魅力
普段はあまり出番がないかもしれないけど、書けると気持ちいい。そして軽い
itertools は、「この処理、for 文でがんばって書いてたけど、実はもっと簡単に書けるのでは？」というときに効いてくる関数がそろっています。
しかも、多くの関数は必要なときに1つずつ値を取り出す仕組み（イテレーター）になっていて、
- 無駄にデータを作らない
- 大きなリストを全部メモリに乗せない
- ループが軽く、処理も速い
という、実はとても効率的な構造になっています。
※イテレーターは「値を1つずつ取り出せる仕組み」のことです。
yield を使ってイテレーターを作る関数のことを「ジェネレーター関数」と呼びます。
yieldやジェネレーターについては以下記事に記載してます！
https://qiita.com/ys-B/items/19511fa0b73e4c0cbfc8
遅延評価
itertools の多くの関数は「遅延評価（lazy evaluation）」になっています。
これは、先述した必要になったときにだけ値を1つずつ取り出す仕組みのことです。
たとえば product() などは、呼び出した瞬間にすべてのデータを作るのではなく、
for 文などで使われたとき、順番に1件ずつ生成します。
たとえば、以下のように product() を使ってみます：
from itertools import product
pairs = product(['A', 'B'], ['X', 'Y'])
print(next(pairs)) # → ('A', 'X')
print(next(pairs)) # → ('A', 'Y')
productメソッドが実行された時点では、すべての組み合わせはまだ作られていません。
next() を呼ぶたびに、1つずつ「そのときに生成される」ようになっています。
同じことを内包表記で書いた場合は…
pairs = [(a, b) for a in ['A', 'B'] for b in ['X', 'Y']]
print(pairs) # → [('A', 'X'), ('A', 'Y'), ('B', 'X'), ('B', 'Y')]
全4パターンが作られてメモリに乗っています。
これがitertools の「遅延評価」です。
必要な分だけ、必要なときに、1つずつ処理される。
だから軽く速いです。
(ちなみに、Python を使っている人にはおなじみの Django の QuerySet も、実は同じように「遅延評価」になっています。)
itertoolsの便利な関数
| 関数名 | なにができる？ | 使用例 |
|---|---|---|
product() |
すべての組み合わせを出す | サイズ × カラーの全組み合わせなど |
combinations() |
リストからn個ずつの組み合わせを出す | テストケース、ペアの列挙など |
compress() |
フラグで要素をフィルタリング | 条件に合う要素だけ抽出したいとき |
count() |
無限に連番を出す | 自動採番、一時的なID付けなど |
groupby() |
ソート済みデータをグループ化して処理 | カテゴリごとの集計・分類など |
おわりに
itertools は、ふだんのコードではなかなか出番がないかもしれません。
でも、ネストした for 文や条件付きのフィルター処理を書いていて
「なんかもっとスッキリ書けないかな……」と思ったとき、
ふと itertools を思い出せると、書くのが楽になって、処理も軽くなる。そんな存在です。
標準ライブラリなので、追加インストールも不要。
まずは product() や compress() あたりから、気軽に試してみてください！
from itertools import *
# ↑ 全部読み込んで、手元で遊んでみるのもおすすめです。