練習問題

以下は permutations 関数を使って解ける問題です。

• ABC054 C - One-stroke Path
• JOI2009予選 D - カード並べ

combinations

combinations は組み合わせの全列挙を行うための関数です。
次のプログラムは [1,2,3,4] というリストから 2 個の要素を取り出す方法を全て出力しています。

コード

import itertools

for p in itertools.combinations([1, 2, 3, 4], 2):
    print(p)

出力

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

combinations 関数は、「タプルを返すジェネレータ」を返します。combinations 関数の戻り値であるジェネレータは、引数に渡したオブジェクトの長さを N 、2 番目の引数として渡した値を K として、必ず {} _ NC _ {K} ²個の要素を持ちます。

特に、第 1 引数に渡したリストが重複する値を含む場合、同じ組み合わせが複数回現れます。

コード

import itertools

for p in itertools.combinations("bbaa", 2):
    print(p)

出力

('b', 'b')
('b', 'a')
('b', 'a')
('b', 'a')
('b', 'a')
('a', 'a')

引数には list, tuple, str, range などのイテラブルオブジェクトを渡すことができます。

combinations_with_replacement

combinations_with_replacement は combinations と同様に組み合わせの全列挙を行うための関数です。
次のプログラムは [1,2,3,4] というリストから重複を許して 2 個の要素を取り出す方法を全て出力しています。

コード

import itertools

for p in itertools.combinations_with_replacement([1, 2, 3], 2):
    print(p)

出力

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

combinations_with_replacement 関数は、「タプルを返すジェネレータ」を返します。combinations_with_replacement 関数の戻り値であるジェネレータは、引数に渡したオブジェクトの長さを N 、2 番目の引数として渡した値を K として、必ず {} _ NH _ {K} 個³の要素を持ちます。

特に、第 1 引数に渡したリストが重複する値を含む場合、同じ組み合わせが複数回現れます。

コード

import itertools

for p in itertools.combinations_with_replacement("bab", 2):
    print(p)

出力

('b', 'b')
('b', 'a')
('b', 'b')
('a', 'a')
('a', 'b')
('b', 'b')

引数には list, tuple, str, range などのイテラブルオブジェクトを渡すことができます。

計算量について

permutations, combinations, combinations_with_replacement の 3 つの関数について、その戻り値を for 文で全てループするときの時間量は O(タプルの長さ \times ループ数) となります。⁴
したがって、通常の for 文と同様の計算量だと考えて問題ありません。参考としていくつかの例を挙げておきます。

product

product 関数は、入れ子になったループを1つにまとめるのに役立ちます。
例えば次のような例を考えます。

コード

for x in range(2):
    for y in range(3):
        print(x, y)

出力

0 0
0 1
0 2
1 0
1 1
1 2

この例では for 文が二重になっています。これを product 関数により次のように書き換えることができます。

import itertools

for x, y in itertools.product(range(2), range(3)):
    print(x, y)

これにより複雑なDPを行う際などにループのネストが深くなることを避けることができます。
product 関数には 3 つ以上の引数を与えることができるため、三重ループ以上も同様に1つのループで書くことができます。

product 関数にはキーワード引数 repeat があり、同じ要素のみを繰り返し引数として渡す以下のケースで記述をまとめることができます。

# 以下の2つは同じになる
itertools.product(range(5), range(5), range(5))
itertools.product(range(5), repeat=3)

これにより、bit DP を itertools.product(range(2),repeat=20) を用いて実装することもできます。

公式ドキュメント

itertools モジュールに含まれる関数のうち、競技プログラミングで使う機会が多いものは以上となりますが、itertools モジュールはここで紹介したもの以外にも様々な関数を含みます。他の関数について知りたい方は公式ドキュメントを参照してください。

問題

本節にはEX問題はありません。本文中にある練習問題を解いてください。

前のページ | 次のページ