Python計算関係ライブラリチートシート ~itertools編~

Why？

競技プログラミングで全探索をする際に知らないと解く手がかりを得られなかったりするので、これを機会にPandasやNumpyなども少しずつ押さえておこうと思ったからです。
とりあえず、今回は座標に関する問題が出たのでitertoolsから以下を抜粋。

itertools.count()

書いたコード

import itertools

# count()で第1引数に指定した数から第2引数で指定した数ずつ返すことができる。
list1 = []
for i in itertools.count(0,2):
    list.append(i)
    if i == 20:
        break
print(list1) # [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]

# 少数を設定することもできる
list2 = []
for i in itertools.count(0,2.5):
    list2.append(i)
    if i == 20:
        break
print(list2) # [0, 2.5, 5.0, 7.5, 10.0, 12.5, 15.0, 17.5, 20.0]

itertools.cycle()

書いたコード

import itertools

# cycle()は引数で指定したiterableを無限にコピーし続ける。
sample = []
count = 0
text = 'TEST'
for i in itertools.cycle(text):
    sample.append(i)
    count += 1
    # 例えばcount == 3だとsample = ['T', 'E', 'S']となる
    if count == len(text)*4:
        break
print(sample) # ['T', 'E', 'S', 'T', 'T', 'E', 'S', 'T', 'T', 'E', 'S', 'T', 'T', 'E', 'S', 'T']

itertools.accumulate()

書いたコード

import itertools

# accumulate()は引数で指定したiterableについて以下のような処理を行う
sample = range(1,6)

result = list(itertools.accumulate(sample))
print(result) # [1, 3, 6, 10, 15]

# つまり[sample[0], sample[0]]+sample[1], sample[0]]+sample[1]+sample[2] .......]ということ

itertools.chain()

書いたコード

import itertools

# chain()は引数で指定したiterableを結合させることができる
sample = 'ABC'
sample2 = 'DEF'

result = list(itertools.chain(sample, sample2))
# join()と違って一文字ずつに格納
print(result) # ['A', 'B', 'C', 'D', 'E', 'F']

# リストに対して使いたい場合はこちら

sample3 = ['ABC', 'DEF']

result = list(itertools.chain.from_iterable(sample3))
print(result) # ['A', 'B', 'C', 'D', 'E', 'F']

itertools.combinations()

書いたコード

import itertools

# combinations()は引数で指定したiterableを第2引数で指定した数での組み合わせを返してくれる

l = [4,4,9,7,5]
new_list = []

# 重複アリの場合はcombinations_with_replacement()と書く
for i in itertools.combinations(l,3):
    new_list.append(i)
print(new_list) # [(4, 4, 9), (4, 4, 7), (4, 4, 5), (4, 9, 7), (4, 9, 5), (4, 7, 5), (4, 9, 7), (4, 9, 5), (4, 7, 5), (9, 7, 5)]

最後に

また問題を問きながらPandas,Numpy含めどんどんアウトプットしていきたいと思います。

参考

公式ドキュメントと以下の記事が簡潔にまとまっていてよかったです。

itertools --- 効率的なループ実行のためのイテレータ生成関数
itertools、more-itertoolsの紹介