Pythonのitertoolsでできる全列挙のまとめ

Pythonのitertoolsモジュールには、イテレーションに関する便利関数が多数用意されています。この記事では、その中でも競技プログラミングで全列挙に使える関数についてまとめます。Python 2.x, 3.xのどちらでも使えます。

Google Code JamのSmallでは全列挙するだけで解ける問題が出ることが多いです。一通りどんな関数があるか知っておきましょう。

import文

スクリプトの冒頭に import itertoolsと書いてあるものとします。

1, 2, 3と書かれた玉が入った袋から、2つの玉を取り出す (3P2の順列)

nums = [1, 2, 3]
for balls in itertools.permutations(nums, 2):
    print(balls)

結果は以下です。

(1, 2)
(1, 3)
(2, 1)
(2, 3)
(3, 1)
(3, 2)

ここで、第2引数を省略すると3P3 (= 3!)が列挙されます。C++のstd::next_permutation()でできることと同じです。

nums = [1, 2, 3]
for balls in itertools.permutations(nums):
    print(balls)

結果は以下です。

(1, 2, 3)
(1, 3, 2)
(2, 1, 3)
(2, 3, 1)
(3, 1, 2)
(3, 2, 1)

1, 2, 3と書かれた玉が入った袋から、2つの玉を区別なく同時に取り出す (3C2の組み合わせ)

nums = [1, 2, 3]
for balls in itertools.combinations(nums, 2):
    print(balls)

結果は以下です。

(1, 2)
(1, 3)
(2, 3)

1, 2, 3と書かれた玉が入った袋から、2つの玉を取り出す。ただし取り出した玉はまた袋に戻す (重複順列)

nums = [1, 2, 3]
for balls in itertools.product(nums, repeat=2):
    print(balls)

結果は以下です。

(1, 1)
(1, 2)
(1, 3)
(2, 1)
(2, 2)
(2, 3)
(3, 1)
(3, 2)
(3, 3)

1, 2, 3と書かれた玉が2セット入った袋から、2つの玉を区別なく同時に取り出す（重複組合せ）

nums = [1, 2, 3]
for balls in itertools.combinations_with_replacement(nums, 2):
    print(balls)

結果は以下です。

(1, 1)
(1, 2)
(1, 3)
(2, 2)
(2, 3)
(3, 3)

2つの集合の直積をとる

signs = ['a', 'b', 'c']
nums = [1, 2, 3]
for pairs in itertools.product(signs, nums):
    print(pairs)

結果は以下です。

('a', 1)
('a', 2)
('a', 3)
('b', 1)
('b', 2)
('b', 3)
('c', 1)
('c', 2)
('c', 3)