Pythonで競プロをやる時にそこそこ使う書き方（itertools編）

誰向け？

競プロをしていて組み合わせを全パターン列挙したいことはありませんか？
私はあります。

そんな時どうしますか？
for文使いますか？
ネストが深くなります。

そもそも書き方がわかりませんか？
確かに...ワカル...

そんな時にPython標準ライブラリのitertoolsが役に立ちます。
自分でも最近使ってみて便利だと思ったので記事にまとめることにしました。

AtCoder最近始めました。Pythonistaです。
みたいな人のお役に立てればなぁと思います。

何が出来るの

itertoolsを使えば順列、組み合わせ、重複順列、重複組合せのパターンを全列挙できます。

import itertools

順列

n個の要素からr個選んで一列に並べる時のパターン$nPr$の列挙
rを省略した場合は$nPn$となります。
iterableには配列以外に文字列とかも入れることが出来ます。

permutations.py

itertools.permutations(iterable, r)

list(itertools.permutations([1, 2, 3, 4], 2))
# -> [(1, 2), (1, 3), (1, 4), (2, 1), (2, 3), (2, 4), 
#     (3, 1), (3, 2), (3, 4), (4, 1), (4, 2), (4, 3)]

list(itertools.permutations([1, 2, 3]))
# -> [(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]

list(itertools.permutations(range(3), 2))
# -> [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)]

組み合わせ

n個の要素からr個選ぶ組み合わせ$nCr$の列挙
permutationsと違って第2引数が必須です。

combinations.py

itertools.combinations(iterable, r)

list(itertools.combinations([1, 2, 3, 4], 2))
# -> [(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]

list(itertools.combinations(range(3), 2))
# -> [(0, 1), (0, 2), (1, 2)]

重複順列

n個の要素から重複を許していち列にr個並べる時のパターン$n^r$の列挙
こちらも第2引数は必須かつrepeat=は省略できません。
→すいません...第2引数なし，repeat省略できる書き方ありました

product.py

itertools.product(iterable, repeat=1)

list(itertools.product(range(3), repeat=2))
# -> [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2)]

list(itertools.product([-1, 1], repeat=3))
# -> [(-1, -1, -1), (-1, -1, 1), (-1, 1, -1), (-1, 1, 1), 
#     (1, -1, -1), (1, -1, 1), (1, 1, -1), (1, 1, 1)]

第2引数を必要としない書き方
ある2次元配列が与えられて，その各行から1つずつ要素を取りたい

product_2.py

T = [[1, 2, 3],
     [4, 5, 6],
     [7, 8, 9]]

list(itertools.product(*T))
# -> [(1, 4, 7), (1, 4, 8), (1, 4, 9), (1, 5, 7), (1, 5, 8), (1, 5, 9), (1, 6, 7), (1, 6, 8), 
# (1, 6, 9), (2, 4, 7), (2, 4, 8), (2, 4, 9), (2, 5, 7), (2, 5, 8), (2, 5, 9), (2, 6, 7), (2, 6, 8), 
# (2, 6, 9), (3, 4, 7), (3, 4, 8), (3, 4, 9), (3, 5, 7), (3, 5, 8), (3, 5, 9), (3, 6, 7), (3, 6, 8), (3, 6, 9)]

重複組合せ

n個の要素から重複を許してr個選ぶ組み合わせ$nHr$の列挙
正直使ったこと無いけど一応

combinations_with_replacement

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

それぞれ、返り値がイテレータになっているので上の例ではlist()をつけましたが、
for文内で使う時はlistはいりません。

for a, b, c in itertools.product([-1, 1], repeat=3):

最後に

実装が重い問題とかだと特に、こういったアルゴリズムの本質じゃない所の手間をなるべく減らしたくなりませんか？itertoolsのおかげでそこの所は大分改善されたかなぁと思います。他にもnumpyやcollections, mathとかをAtCoderでは使うので、機会があればまたまとめるかもしれません。最後に、実際に自分がitertoolsを使って解いた問題を置いておきます。

最近itertoolsを使った問題

[ABC028 C - 数を3つ選ぶマン]
(https://beta.atcoder.jp/contests/abc028/tasks/abc028_c)
[ABC029 C - Brute-force Attack]
(https://beta.atcoder.jp/contests/abc029/tasks/abc029_c)
ABC073 D - joisino's travel
ABC100 D - Patisserie ABC
ABC114 C - 755
ABC015 C - 高橋くんのバグ探し