LoginSignup
123
132

More than 3 years have passed since last update.

Pythonで競プロするのに必要な機能をまとめてみた~itertools~

Last updated at Posted at 2020-05-08

初めに

ABC165のC問題が解けず、コンテスト後にツイッターで解法を見て回っていたところ、以下のようなmaspyさんのツイートを拝見し、itertoolsの関数はよく使うものの整理できていないと感じました。
また、itertools関連の記事はQiitaにあるものの、競技プログラミングにおける使用法に触れつつ網羅的に扱っている記事がなかったので、記事としてまとめることにしました。

itertoolsとは??

効率的なループ実行のためのイテレータを生成する関数を集めたモジュールのことです。
イテレータとは複数の値を持つデータ構造のデータに対して順番にアクセスするための機能を与えるもののことです。
ここではitertoolsの関数の機能に注目したいので、イテレータについての詳しい説明は割愛します。

itertoolsの関数

関数名
(1) accumulate関数
(2) chain関数,chain.from_iterable関数
(3) combinations関数
(4) combinations_with_replacement関数
(5) compress関数
(6) dropwhile関数,takewhile関数
(7) groupby関数
(8) islice関数
(9) permutations関数
(10) product関数
(11) starmap関数
(12) zip_longest関数

それぞれの関数は、機能→使用方法とその例→使用可能な問題、の順番で紹介しています。
また、競技プログラミングにおいて自分が有用と感じた関数や機能を抽出しているので、他にもある場合は教えていただければ幸いです。

さらに、以下では、コードにおける表記を省略するために

from itertools import *

の形でitertoolsをimport済みとしているので関数を使用する際は注意してください。

追記(2020/05/11)

以下の例は、Python3.8のインタラクティブモードで確認済みです。
また、pypyでは一部の引数が使えずREが出る可能性があるので、提出の際にはお気をつけください。

(1)accumulate関数

機能

accumulate(iterable[,func=operator.add,initial=None])
'''
funcによるiterableの累積の結果をイテレータで返す関数

引数
----------
iterable:イテラブルなオブジェクト
func:二つの引数を持つ関数
initial:初期値(ver3.8で追加)
'''

使用方法とその例

通常は以下のように累積和に使用します。

>>> iterable=[3, 4, 6, 2, 1, 9, 0, 7, 5, 8]

#通常の累積和
>>> list(accumulate(iterable))
[3, 7, 13, 15, 16, 25, 25, 32, 37, 45]

funcやinitial次第では以下のような使い方もできます。
また、iterableの要素は整数だけではなく、funcの引数として受け入れられる任意の型で用いることができます。
さらに、最終的な累積値(末尾の値)だけを求めたい場合はfunctoolsモジュールのreduce関数を利用して求めることができます。
競技プログラミングでは、累積和,累積最小値,累積最大値は良く使うと思います。

>>> iterable=[3, 4, 6, 2, 1, 9, 0, 7, 5, 8]

#累積最小値
>>> list(accumulate(iterable, func=min))
[3, 3, 3, 2, 1, 1, 0, 0, 0, 0]

#累積最大値
>>> list(accumulate(iterable, func=max))
[3, 4, 6, 6, 6, 9, 9, 9, 9, 9]

#累積積(operatorモジュールのimportが必要)
>>> import operator
>>> list(accumulate(iterable, func=operator.mul))
[3, 12, 72, 144, 144, 1296, 0, 0, 0, 0]

#初期値を設定した場合(累積和)→配列の初めに初期値が追加される
>>> list(accumulate(iterable, initial=100))
[100, 103, 107, 113, 115, 116, 125, 125, 132, 137, 145]

#初期値を設定できない場合(PyPy7.3.0で)


#iterableの要素は文字列でも可
>>> iterable_S=["A", "B", "C", "D", "E"]
>>> list(accumulate(iterable_S))
['A', 'AB', 'ABC', 'ABCD', 'ABCDE']

#reduce関数による累積和
>>> import functools, operator
>>> functools.reduce(operator.add, iterable)
45

使用可能な問題

(まだまだたくさんありますが、ここでは省略します。)
ABC047-D 高橋君と見えざる手
ABC152-C Low Elements
ABC122-C GeT AC
ABC134-C Exception Handling
ARC098-C Attention
AGC023-A Zero-Sum Ranges
ABC182-D D - Wandering
Educational Codeforces 86-C Yet Another Counting Problem
Educational Codeforces 90-C Pluses and Minuses
ABC172-C Tsundoku

(2)chain関数,chain.from_iterable関数

機能

chain(*iterables)
'''
与えられたiterablesの全要素を前のオブジェクトから順番にイテレータで返す関数

引数
*iterables:(複数の)イテラブルなオブジェクト
'''

chain.from_iterable(iterable)
'''
与えられたiterableをアンパックしてchainと同様の操作をする関数

引数
iterable:1個のイテラブルなオブジェクト
'''

使用方法とその例

通常は以下のように複数のイテラブルなオブジェクト(異なる型どうしでもOK)を連結して1個のイテラブルにすることができます。
また、chain関数は複数のイテラブルなオブジェクトを引数で指定するのに対し、chain.from_iterable関数はそれらのオブジェクトを格納した1個のイテラブルなオブジェクトを引数で指定する必要があるので、違いに注意が必要です。

>>> iterable1=(0, 1, 2)
>>> iterable2=[3, 4, 5]

#chain関数を通常通り用いた場合
>>> list(chain(iterable1, iterable2))
[0, 1, 2, 3, 4, 5]

#chain.from_iterable関数を通常通り用いた場合
>>> list(chain.from_iterable([iterable1, iterable2]))
[0, 1, 2, 3, 4, 5]

#chain.from_iterable関数でchain関数のように用いようとした場合
>>> list(chain.from_iterable(iterable1, iterable2))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: from_iterable() takes exactly one argument (2 given)

競技プログラミングでは、以下のように二次元配列での平坦化で使うことができます。また、この際はchain.from_iterable関数を使う必要があることに注意してください。

>>> iterable=[
... [0, 1, 2],
... [3, 2, 4],
... [4, 1, 5]
... ]
>>> list(chain.from_iterable(iterable))
[0, 1, 2, 3, 2, 4, 4, 1, 5]

使用可能な問題

(他にもある場合は有識者の方教えてください。)
ABC126-F XOR Matching
ABC159-E Dividing Chocolate
square869120Contest #6 B - AtCoder Market

(3)combinations関数

機能

combinations(iterable, r)
'''
iterableの要素のみからなる長さrの部分列をタプルで辞書順にイテレータで返す関数
ただし、r>nの場合は何も返さない

引数
----------
iterable:イテラブルなオブジェクト(長さはn)
r:部分列の長さ(int型)
'''

使用方法とその例

以下のように、n個の要素からr個選ぶ組み合わせを全て列挙することができます($_nC_r$通り)。
また、組み合わせは位置に基づいた辞書順で列挙されるので、引数のiterableがソートされている場合は(その要素の)辞書順(昇順)で出力されます(iterable1,iterable4)。
さらに、要素に重複がある場合も、位置のみで異なるかどうかを判定するので出力されるタプルに重複が生じることに注意が必要です(iterable3)。

#整数が昇順で並んでいる場合→イテレータの出力順も昇順
>>> iterable1=(0, 1, 2, 3)
>>> list(combinations(iterable1, 2))
[(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]

#整数がランダムな順番で並んでいる場合→イテレータの出力順は(場所に基づいた)昇順
>>> iterable2=(0, 2, 1, 3)
>>> list(combinations(iterable2, 2))
[(0, 2), (0, 1), (0, 3), (2, 1), (2, 3), (1, 3)]

#整数に重複がある場合→イテレータの出力順は(場所に基づいた)昇順でタプルに重複が生じる
>>> iterable3=(0, 0, 1, 2)
>>> list(combinations(iterable3, 2))
[(0, 0), (0, 1), (0, 2), (0, 1), (0, 2), (1, 2)]

#文字列が辞書順で並んでいる場合→イテレータの出力順も辞書順
>>> iterable4=("A", "B", "C", "D")
>>> list(combinations(iterable4, 2))
[('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('C', 'D')]
>>> ["".join(i) for i in list(combinations(iterable4,2))]
['AB', 'AC', 'AD', 'BC', 'BD', 'CD']

また、L,R,Nを正の整数として、数列$A_1,A_2,…,A_N$について$L \leqq A_1< A_2 < … < A_N \leqq R$が成り立つ時、combinations(range(L, R+1), N+1)でありうる数列$A_1,A_2,…,A_N$をタプルで列挙することができます。

#N=4,L=1,R=4
>>> iterable=range(1, 5)
>>> list(combinations(iterable, 3))
[(1, 2, 3), (1, 2, 4), (1, 3, 4), (2, 3, 4)]

使用可能な問題

(他にもある場合は有識者の方教えてください。)
ABC147-C HonestOrUnkind2

(4)combinations_with_replacement関数

機能

combinations_with_replacement(iterable, r)
'''
iterableの要素のみからなる長さrの部分列(重複を許す)をタプルで辞書順にイテレータで返す関数

引数
----------
iterable:イテラブルなオブジェクト(長さをnとします)
r:部分列の長さ(int型)
'''

使用法とその例

(このまとめ記事を書くきっかけになった関数です。)

以下のように、n個の要素から重複してr個選ぶ組み合わせを全て列挙することができます($_{n+r-1}C_r$通り→導出はこの記事を参照)。
また、組み合わせは位置に基づいた辞書順で列挙されるので、引数のiterableがソートされている場合は(その要素の)辞書順(昇順)で出力されます(iterable1,iterable4)。
さらに、要素に重複がある場合も、位置のみで異なるかどうかを判定するので出力されるタプルに重複が生じることに注意が必要です(iterable3)。

#整数が昇順で並んでいる場合→イテレータの出力順も昇順
>>> iterable1=(0, 1, 2, 3)
>>> list(combinations_with_replacement(iterable1,2))
[(0, 0), (0, 1), (0, 2), (0, 3), (1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]

#整数がランダムな順番で並んでいる場合→イテレータの出力順は(場所に基づいた)昇順
>>> iterable2=(0, 2, 1, 3)
>>> list(combinations_with_replacement(iterable2,2))
[(0, 0), (0, 2), (0, 1), (0, 3), (2, 2), (2, 1), (2, 3), (1, 1), (1, 3), (3, 3)]

#整数に重複がある場合→イテレータの出力順は(場所に基づいた)昇順でタプルに重複が生じる
>>> iterable3=(0, 0, 1, 2)
>>> list(combinations_with_replacement(iterable3,2))
[(0, 0), (0, 0), (0, 1), (0, 2), (0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)]

#文字列が辞書順で並んでいる場合→イテレータの出力順も辞書順
>>> iterable4=("A", "B", "C", "D")
>>> list(combinations_with_replacement(iterable4,2))
[('A', 'A'), ('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'B'), ('B', 'C'), ('B', 'D'), ('C', 'C'), ('C', 'D'), ('D', 'D')]
>>> ["".join(i) for i in list(combinations_with_replacement(iterable4,2))]
['AA', 'AB', 'AC', 'AD', 'BB', 'BC', 'BD', 'CC', 'CD', 'DD']

また、L,R,Nを正の整数として、数列$A_1,A_2,…,A_N$について$L \leqq A_1 \leqq A_2 \leqq … \leqq A_N \leqq R$が成り立つ時、combinations_with_replacement(range(L,R+1),N+1)でありうる数列$A_1,A_2,…,A_N$をタプルで列挙することができます。

#N=4,L=1,R=4
>>> iterable=range(1, 5)
>>> list(combinations_with_replacement(iterable, 3))
[(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 1, 4), (1, 2, 2), (1, 2, 3), (1, 2, 4), (1, 3, 3), (1, 3, 4), (1, 4, 4), (2, 2, 2), (2, 2, 3), (2, 2, 4), (2, 3, 3), (2, 3, 4), (2, 4, 4), (3, 3, 3), (3, 3, 4), (3, 4, 4), (4, 4, 4)]

使用可能な問題

(他にもある場合は有識者の方教えてください。)
ABC165-C Many Requirements

(5)compress関数

機能

compress(iterable, selectors)
'''
iterableの要素に対応するselectorsの要素がTrueと評価されるもののみをイテレータで返す関数

引数
----------
iterable:イテラブルなオブジェクト
selectors:イテラブルなオブジェクト
'''

使用法とその例

以下のように、フィルタをかけて特定の要素を取り出すのに使います。
同様の機能を持つ関数として組み込みのfilter関数がありますが、compress関数はiterableとselectorsのうち短い方に長さが合わせられます。
また、filter関数は条件によるフィルター、compress関数は位置によるフィルターをしているという違いがあります。
さらに、リスト内包表記を使用するとfilter関数およびcompress関数と同様のことをすることができます。

>>> iterable=[1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89] #長さ11
>>> selectors1=[0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0] #長さ12
>>> selectors2=[0, 0, 1, 0, 0, 1, 0, 0] #長さ8
>>> f=lambda x:x%2==0

#selectorsの長さがiterableの長さ(以上)の場合→iterableの要素に対応するselectorsの要素が1(True)であれば取り出される(最大の返される要素数はiterableの長さ)
>>> list(compress(iterable,selectors1))
[2, 8, 34]

#selectorsの長さがiterableの長さ未満の場合→1(True)であるselectorsの要素に対応するiterableの要素が取り出される(最大の返される要素数はselectorsの長さ)
>>> list(compress(iterable,selectors2))
[2, 8]

#fの返り値がtrueになる要素(ここでは偶数の要素)のみがiterableから取り出される
>>> list(filter(f,iterable))
[2, 8, 34]

#上記と同様のことをリスト内包表記でもできる
>>> [num for i,num in enumerate(iterable) if selectors1[i]]
[2, 8, 34]
>>> [num for i,num in enumerate(iterable) if i<ls and selectors2[i]]
[2, 8]
>>> [num for num in iterable if f(num)]
[2, 8, 34]

使用可能な問題

compress関数を使うと便利になる問題があれば有識者の方教えてください。

(6)dropwhile関数,takewhile関数

機能

dropwhile(prediictable, iterable)
'''
predictableが真である場合は要素を飛ばし、(初めに)偽である要素以降のイテレータを返す関数

引数
----------
predictable:iterableの要素を引数として持つ関数
iterable:イテラブルなオブジェクト
'''

takewhile(prediictable, iterable)
'''
predictableが真である場合はその要素のイテレータを返し、(初めに)偽である要素以降は無視する関数

引数
----------
predictable:iterableの要素を引数として持つ関数
iterable:イテラブルなオブジェクト
'''

使用法とその例

上記を見ればわかるようにdropwhile関数とtakewhile関数はちょうど逆の関係にあります。
以下のようにx<5を(初めに)満たさない要素(インデックス3)を境にdropwhile関数の戻り値とtakewhile関数の戻り値が異なることがわかります。

>>> iterable=[1, 2, 3, 5, 1, 4, 2, 6]

#dropwhile関数はx<5を(初めに)満たさない要素(インデックス3)以降の要素を返す。
>>> list(dropwhile(lambda x:x<5,iterable))
[5, 1, 4, 2, 6]

#takewhile関数はx<5を(初めに)満たさない要素(インデックス3)より前の要素を返す。
>>> list(takewhile(lambda x:x<5,iterable))
[1, 2, 3]

また、競技プログラミングでの具体的な使用方法については、以下のようなものがあります。
以下を見ればわかるようにbreakするポイントの前後のイテラブルなオブジェクトが欲しい場合には強力だと考えられます。

#√nよりも小さい要素だけ取り出す問題を考える
>>> n=100

#√nの計算に誤差が出る可能性がある→takewhile関数を使いたい!
>>> list(range(int(n**0.5)))
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

#2乗がnよりも小さい要素だけ取り出す
>>> list(takewhile(lambda x:x*x<n,range(n)))
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

#リスト内包表記でも書ける
>>> [i for i in range(n) if i*i<n]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

使用可能な問題

dropwhile関数,takewhile関数を使うと便利になる問題があれば有識者の方教えてください。
yukicoder No.1219 Mancala Combo
ABC105-C Base -2 Number
ABC167-F Bracket Sequencing

(7)groupby関数

機能

gropuby(iterable, key=None)
'''
iterable中の同じkeyを持つ要素からなる(連続した)グループについて、そのキーとグループを組にしてイテレータで返す関数

引数
----------
iterable:イテラブルなオブジェクト
key:関数を指定(指定しない場合は恒等関数)
'''

使用法とその例

以下のように、同じkeyを持つ要素をiterableの前の要素から順にグループ化することができます。
この際の戻り値のグループについてはitertools._grouperオブジェクトになっているのでリストなどに変換してから使う必要があります。

>>> iterable=[1, 1, 2, 2, 2, 4, 1, 3, 5]

#itertools._grouperオブジェクトでグループは返される
>>> [(key,group) for key,group in groupby(iterable)]
[(1, <itertools._grouper object at 0x106bb53c8>), (2, <itertools._grouper object at 0x106bb52b0>), (4, <itertools._grouper object at 0x106bb52e8>), (1, <itertools._grouper object at 0x106bb5240>), (3, <itertools._grouper object at 0x106bb5278>), (5, <itertools._grouper object at 0x106bb54e0>)]

#グループをリストへと変換する(この場合は同じ要素どうしがグループ化)
>>> [(key,list(group)) for key,group in groupby(iterable)]
[(1, [1, 1]), (2, [2, 2, 2]), (4, [4]), (1, [1]), (3, [3]), (5, [5])]

また、ソートしてからgorupby関数を適用すると、連続してない要素もまとめてグループ化することができ、似たことがcollectionsモジュールのCounterクラスによってできます。ただし、Counterクラスではkeyについて整列されていないので注意が必要です。
さらに、keyに何も指定しない場合は恒等関数ですが、一番最後の例のようにlambda x: x%2を指定すれば偶奇によるグループ分けなどもできます。

groupby関数の有用さがこの記事のみでは伝わりにくいので、競技プログラミングにおける使い方を競プロの応用事項確認~ランレングス圧縮~にて紹介しているので、ぜひ読んでいただきたいです。

>>> iterable=[1, 1, 2, 2, 2, 4, 1, 3, 5]

#連続してない要素もまとめてグループ化することができる
>>> [(key,list(group)) for key,group in groupby(sorted(iterable))]
[(1, [1, 1, 1]), (2, [2, 2, 2]), (3, [3]), (4, [4]), (5, [5])]

#グループの長さだけを組にすることもできる(①)
>>> [(key,len(list(group))) for key,group in groupby(sorted(iterable))]
[(1, 3), (2, 3), (3, 1), (4, 1), (5, 1)]

#Counterはdictのサブクラスなので、dictと同様の扱いができる
#iterableの要素中にそれぞれの要素がいくつあるかをカウントすることができる
>>> from collections import Counter
>>> Counter(iterable)
Counter({1: 3, 2: 3, 4: 1, 3: 1, 5: 1})

#①と同じ形にすることができる
>>> sorted(Counter(iterable).items(), key=lambda x:x[0])
[(1, 3), (2, 3), (3, 1), (4, 1), (5, 1)]

#keyに指定する関数次第では、偶奇でグループ分けすることなどもできる
>>> [(key,list(group)) for key,group in groupby(iterable,key=lambda x: x%2)]
[(1, [1, 1]), (0, [2, 2, 2, 4]), (1, [1, 3, 5])]

使用可能な問題

ABC019B 高橋くんと文字列圧縮
ABC124D Handstand
ABC047C 一次元リバーシ
ABC182E Akari
HHKB2020E Lamps
ABC043D アンバランス
Codeforces #603(Div.2)A Sweet Problem
Codeforces #481(Div.3)B File Name
M-SOLUTIONS D Road to Millionaire
AGC040A ><
ABC136D Gathering Children

(8)islice関数

機能

islice(iterable[, start], stop[, step])

'''
iterableから要素を選択(スライス)してイテレータで返す関数

引数
----------
iterable:イテラブルなオブジェクト
start:start(インデックス)から要素を選択する(Noneまたは省略する時はstart=0と同じ)
stop:stop(インデックス)の一つ手前まで要素を選択する(省略する事はできず,Noneの時はstop=(オブジェクトの長さ)と同じ)
step:いくつ飛ばしで要素を選択するかを指定する(Noneまたは省略する時はstep=1と同じ,使用する際はstart引数も必要)
'''

使用法とその例

リストや文字列やタプルにある[start:stop:step]を任意のイテラブルなオブジェクトについて用いることができるようにしたのがこの関数です。
isliceは便利で使えるタイミングは多いですが、競技プログラミングに限れば[start:stop:step]を使えれば良いと思います。

>>> iterable=range(10)
#以下では、記述がない場合、start=0,stop=10-1=9,step=1

#start=3,stop=6の場合,list(iterable[3:6])と同じ
>>> list(islice(iterable,3,6))
[3, 4, 5]

#stop=3の場合,list(iterable[:3])と同じ
>>> list(islice(iterable,3))
[0, 1, 2]

#start=None,stop=3の場合,list(iterable[:3])と同じ(Noneは省略可能)
>>> list(islice(iterable,None,3))
[0, 1, 2]

#start=3,stop=Noneの場合,list(iterable[3:])と同じ(Noneは省略不可能!!)
>>> list(islice(iterable,3,None))
[3, 4, 5, 6, 7, 8, 9]

#start=3,stop=None,step=2の場合,list(iterable[3:10:2])と同じ
>>> list(islice(iterable,3,None,2))
[3, 5, 7, 9]

使用可能な問題

(islice関数を使うと便利になる問題があれば有識者の方教えてください。)

(9)permutations関数

機能

permutations(iterable, r=n])
'''
iterableの要素のみからなる長さrの順列をタプルで辞書順にイテレータで返す関数
ただし、r>nの場合は何も返さない

引数
----------
iterable:イテラブルなオブジェクト(長さをnとします)
r:順列の長さ(int型),指定しない場合は全順列(r=n)になる
'''

使用法とその例

以下のように、n個の要素からr個を並べる順列を全て列挙することができます($_nP_r$通り)。
また、順列は位置に基づいた辞書順で列挙されるので、引数のイテラブルなオブジェクトがソートされている場合は(その要素の)辞書順(昇順)で出力されます。
さらに、要素に重複がある場合も、位置のみで異なるかどうかを判定するので出力されるタプルに重複が生じることに注意が必要です。

#整数が昇順で並んでいる場合→イテレータの出力順も昇順
>>> iterable1=(0, 1, 2, 3)
>>> list(permutations(iterable1, 2))
[(0, 1), (0, 2), (0, 3), (1, 0), (1, 2), (1, 3), (2, 0), (2, 1), (2, 3), (3, 0), (3, 1), (3, 2)]

#整数がランダムな順番で並んでいる場合→イテレータの出力順は(場所に基づいた)昇順
>>> iterable2=(0, 2, 1, 3)
>>> list(permutations(iterable2, 2))
[(0, 2), (0, 1), (0, 3), (2, 0), (2, 1), (2, 3), (1, 0), (1, 2), (1, 3), (3, 0), (3, 2), (3, 1)]

#整数に重複がある場合→イテレータの出力順は(場所に基づいた)昇順でタプルに重複が生じる
>>> iterable3=(0, 0, 1, 2)
>>> list(permutations(iterable3, 2))
[(0, 0), (0, 1), (0, 2), (0, 0), (0, 1), (0, 2), (1, 0), (1, 0), (1, 2), (2, 0), (2, 0), (2, 1)]

#文字列が辞書順で並んでいる場合→イテレータの出力順も辞書順
>>> iterable4=("A", "B", "C", "D")
>>> list(permutations(iterable4, 2))
[('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'A'), ('B', 'C'), ('B', 'D'), ('C', 'A'), ('C', 'B'), ('C', 'D'), ('D', 'A'), ('D', 'B'), ('D', 'C')]
>>> ["".join(i) for i in list(permutations(iterable4, 2))]
['AB', 'AC', 'AD', 'BA', 'BC', 'BD', 'CA', 'CB', 'CD', 'DA', 'DB', 'DC']

使用可能な問題

(他にもある場合は有識者の方教えてください。)
ABC073-D joisino's travel
ABC054-C One-stroke Path
Codeforces #592(Div.2)-D Paint the Tree
ABC181D - Hachi

(10)product関数

機能

product(*iterables[, repeat=1])
'''
iterablesに含まれるイテラブルなオブジェクトの直積(デカルト積)の要素を順にイテレータで生成する関数
repeatに指定された数の分だけイテラブルなオブジェクトは繰り返される。

引数
----------
*iterables:(複数の)イテラブルなオブジェクト
repeat:iterablesを繰り返す回数(int型)
'''

使用法とその例

以下のように、直積の要素を全列挙することができます。
また、直積の要素は位置に基づいた辞書順で列挙されるので、引数のイテラブルなオブジェクトがソートされている場合は(その要素の)辞書順(昇順)で出力されます(✳︎)。
さらに、要素に重複がある場合も、位置のみで異なるかどうかを判定するので出力されるタプルに重複が生じることに注意が必要です(✳︎)。
加えて、repeatにkを設定するとiterablesがk回繰り返されて直積が生成されます。例えば、product(x,y,repeat=2)の場合はproduct(x,y,x,y)と同様の生成結果を返します。

(✳︎)…これらのパターンの例はここでは省くので、自分で確かめてみてください。

>>> iterable1=[0, 1]
>>> iterable2=[2, 3]

#iterable1とiterable2の直積を生成する
>>> list(product(iterable1, iterable2))
[(0, 2), (0, 3), (1, 2), (1, 3)]

#3個のiterable1の直積
>>> list(product(iterable1, 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)]

#2個のiterable1とiterable2の直積
>>> list(product(iterable1, iterable2, repeat=2))
[(0, 2, 0, 2), (0, 2, 0, 3), (0, 2, 1, 2), (0, 2, 1, 3), (0, 3, 0, 2), (0, 3, 0, 3), (0, 3, 1, 2), (0, 3, 1, 3), (1, 2, 0, 2), (1, 2, 0, 3), (1, 2, 1, 2), (1, 2, 1, 3), (1, 3, 0, 2), (1, 3, 0, 3), (1, 3, 1, 2), (1, 3, 1, 3)]

#「2個のiterable1とiterable2の直積」は「iterable1とiterable2とiterable1とiterable2の直積」に等しい
>>> list(product(iterable1, iterable2, repeat=2)) == list(product(iterable1, iterable2, iterable1, iterable2))
True

また、競技プログラミングではn進数の生成やn進数のbit全探索で強力な効果を発揮します。
以下は、3進数で桁数が3以下のものを生成しています。
イメージが湧きにくいかもしれないので、使用可能な問題で試してください。

>>> ternary=range(3)

#以下、3進数が小さい順に生成されていることがわかる
#3進数のbit全探索もこれを使えば行うことができる
>>> list(product(ternary,repeat=3))
[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 1, 0), (0, 1, 1), (0, 1, 2), (0, 2, 0), (0, 2, 1), (0, 2, 2), (1, 0, 0), (1, 0, 1), (1, 0, 2), (1, 1, 0), (1, 1, 1), (1, 1, 2), (1, 2, 0), (1, 2, 1), (1, 2, 2), (2, 0, 0), (2, 0, 1), (2, 0, 2), (2, 1, 0), (2, 1, 1), (2, 1, 2), (2, 2, 0), (2, 2, 1), (2, 2, 2)]

使用可能な問題

ABC031-D 語呂合わせ
ABC114-C 755
ABC119-C Synthetic Kadomatsu
ABC162-C Sum of gcd of Tuples (Easy)

(11)starmap関数

機能

starmap(function,iterable)
'''
iterableの要素を引数としてfunctionを計算し、その結果をイテレータを生成して返す関数

引数
----------
function:iterableの要素を変換する関数
iterable:イテラブルなオブジェクト
'''

使用法とその例

競プロでよく使う組み込みのmap関数と似た機能を持ちます。
しかし、map関数ではiterableの要素がfunctionにそのまま引数として与えられるのに対し、starmap関数ではiterableの要素はグループ化されており(つまり、その要素もイテラブル)、グループ化された要素に含まれるそれぞれの要素が引数として与えられるので、注意が必要です。

>>> iterable=[[0,1,1,2],[2,3,4,4],[5,1,2,6]]

#iterableの要素である配列はsumにそのまま引数として与えることができる
>>> list(map(sum,iterable))
[4, 13, 14]

#iterableのそれぞれのイテラブルな要素(配列)の長さは4なので、functionの引数も4個必要
>>> list(starmap(sum,iterable))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: sum expected at most 2 arguments, got 4

#引数を4つにすれば、iterableの要素である配列のそれぞれの要素が引数としてfunctionに与えられるので、starmapでも使用可能
>>> list(starmap(lambda a,b,c,d:a+b+c+d,iterable))
[4, 13, 14]

#functionの引数の数がiterableの要素の配列の長さと異なると使うことはできない
>>> list(starmap(lambda a,b,c:a+b+c,iterable))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: <lambda>() takes 3 positional arguments but 4 were given

>>> iterable_S=["1","2","3"]

#iterableの要素である整数文字列はintにより整数に変換することができる
>>> list(map(int,iterable_S))
[1, 2, 3]

#iterableのそれぞれのイテラブルな要素(文字列)の長さは1なので、functionの引数は1個で良い
#→この場合、mapとstarmapの実行結果は見かけ上同じになる
>>> list(starmap(int,iterable_S))
[1, 2, 3]

競技プログラミングでは組み込みのzip関数と組み合わせれば使える機会があるかもしれません。
しかし、zipされたオブジェクトをunzipすればmap関数(iterableを引数として複数与える)でも同様のことができるので、map関数を使えれば困ることはなさそうです。

>>> x=[0,0,1,2,3]
>>> y=[5,1,4,2,3]

#zip関数でx,yの要素を順にタプルで組にしているので、zip(x, y)の要素は長さ2のタプル(→functionの引数は2個)
>>> list(starmap(lambda a, b: a+b, zip(x, y)))
[5, 1, 5, 4, 6]

#mapの場合はタプルとして与えるので、引数は1個
>>> list(map(lambda x: x[0]+x[1], zip(x, y)))
[5, 1, 5, 4, 6]

#unzip(zip(x,y)→x,y)すると、(引数が複数でも)mapの引数にiterableの要素はそのまま与えられる
>>> list(map(lambda a, b: a+b, x, y))
[5, 1, 5, 4, 6]

使用可能な問題

(starmap関数を使うと便利になる問題があれば有識者の方教えてください。)

(12)zip_longest関数

機能

zip_longest(*iterables[ ,fillvalue=None])
'''
各iterableの要素を順にタプルとしてまとめる関数
iterableの要素の長さが違う場合はfillvalueによって埋められる

引数
----------
*iterables:(複数の)イテラブルなオブジェクト
fillvalue:iterableの長さの違いを補ってタプルの長さを統一する要素
'''

使用法とその例

組み込みのzip関数と似た機能を持ちます。
この二つの関数では、zip関数がタプルとしてまとめる際に最も長さが短いiterableに長さを合わせるのに対し、zip_longest関数はタプルとしてまとめる際に最も長さが長いiterableに長さを合わせ、足りない部分はfillvalueで埋める、という違いがあります。
また、競技プログラミングでは、for文を回してxとyの要素(同じ長さのイテラブルなオブジェクト)を同時に取ってきたい時に、for x_sub, y_sub in zip(x, y):と書くと便利です。

>>> x=[0,1,2,3,4]
>>> y=[5,6,7,8,9,10]

#zip関数では長さの短いxに合わせてタプルで要素がまとめられる
>>> list(zip(x, y))
[(0, 5), (1, 6), (2, 7), (3, 8), (4, 9)]

#zip_longest関数では長さの長いyに合わせてタプルで要素がまとめられる(fillvalue=None)
>>> list(zip_longest(x, y))
[(0, 5), (1, 6), (2, 7), (3, 8), (4, 9), (None, 10)]

#fillvalueに-1を指定すれば以下のようになる
>>> list(zip_longest(x, y, fillvalue=-1))
[(0, 5), (1, 6), (2, 7), (3, 8), (4, 9), (-1, 10)]

使用可能な問題

(zip_longest関数を使うと便利になる問題があれば有識者の方教えてください。)

省略したitertoolsの関数

以下は競技プログラミングでは使いにくいと感じた関数なので省略しました。

count関数
cycle関数
filterfalse関数
repeat関数
tee関数

参考ページ

以下は参考にさせていただいたページの一覧になります(2020/05/09更新)。

itertools全般
Pythonのドキュメント(itertools)
すごいぞitertoolsくん
Pythonで競プロやるときによく書くコードをまとめてみた
Python3のitertoolsモジュールをまとめてみました!
Pythonでこんなことできちゃうんです
Pythonメモ-72 (ライブラリメモ - itertools)

(1)accumulate関数
累積和とは(超初心者用)

(2)chain関数,chain.from_iterable関数
chain, chain.from_iterableの紹介(pythonのitertoolsを使いこなすために)

(8)islice関数
isliceの紹介、具体例添え(pythonのitertoolsを使いこなすために)

(10)product関数
Pythonで複数のリストの直積(デカルト積)を生成するitertools.product

(11)starmap関数
【python】mapで複数の引数を渡したいときはstarmapが便利

(12)zip_longest関数
Python, zip関数の使い方: 複数のリストの要素をまとめて取得

123
132
1

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
123
132