なに
Atcoderなどのプログラミングコンテストでいろいろ試行錯誤しているうちに、itertoolsというライブラリにたどり着いた。古来このライブラリの使い勝手の良さには多くの人が感動してきたと思うが、私もご他聞に漏れず感激したので、いつでも使えるように記録しておく。
いろいろできるぞ
累積和
ary = [1, 3, 5, 7, 9]
cumsum = itertools.accumulate(ary)
accumulateを使うとリストの累積和のイテレータが返ってくる。すごい。
print(list(cumsum))
# -> [1, 4, 9, 16, 25]
イテレータなので一気に出力するにはlistを使う。numpy.cumsumでももちろん可能だけど……
s = ['ab', 'bc', 'cd']
print(list(itertools.accumulate(s)))
# -> ['ab', 'abbc', 'abbccd']
文字列でも使えるのはすごい。
ある値が何個あるか
groupbyという関数には大変感動した。
bi = [0,0,0,1,1,0,0,0,1,1,0,1]
こんなリストにgroupbyを使うと、値をkey,同じ値の集まりをgroupとしたイテレータを返してくる。
gr = itertools.groupby(bi)
for key, group in gr:
print(f'{key}: {list(group)}')
# 0: [0, 0, 0]
# 1: [1, 1]
# 0: [0, 0, 0]
# 1: [1, 1]
# 0: [0]
# 1: [1]
というように、リストの最初から数えて同じ値のものを区切って出力してくれる。groupがイテレータになっているのもあとの処理がいろいろできて楽しい。Atcoderでは大変お世話になりました。
忘れちゃいけない組み合わせ
itertoolsの真骨頂ともいえる組み合わせの出力。簡単な全探索ならこれで出力しちゃってもいいんじゃないだろうか(速度はどうなんだろう)。
順列
permutationsという関数にiterableなものを渡すと、全部の順列組み合わせを出してくれる。
print(list(itertools.permutations([1, 2, 3])))
# -> [(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]
第2引数には何個取り出すか選べる。
文字列もiterableなのでもちろん可。
print(list(itertools.permutations('Qiita')))
# ->冗談でやったらめっちゃ出力出てきたんで省略。まあ120通りだもんね……。
組み合わせ
もちろんCombinationsもある。
print(list(itertools.combinations([1, 2, 3], 2)))
# -> [(1, 2), (1, 3), (2, 3)]
調べて知ったけど、重複を許すCombinationsもあるらしい。
print(list(itertools.combinations_with_replacement([1, 2, 3], 2)))
# -> [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]
ちょっと使いどころがわからなかった。
直積も便利。例えば0と1の組み合わせをrepeatで指定した桁出力することができる。
print(list(itertools.product([0,1], repeat=3)))
# -> [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1)]
まあビットで考えたら早いんだけど、リストにできるのはちょっぴり嬉しい。
結局は
ここにすべてが載っている……。わかりやすい。
やりたいこと
上でやれることはだいたいnumpyとかでもできるはずで(cumsumとか)、どちらが速いのかはちょっと気になる。itertoolsのほうが遅いみたいな記事を見たけど、数字は書いていなかったし、都度生成のitertoolsとそうではないnumpy.arrayを比較するにはちょっと工夫がいる気がする。
そんなことを出張中の新幹線車内で書きなぐっていたのでした。