【Python】順列全探索入門！リストの要素を並べ替えて答えを探す

リストの要素を並び替えて、すべての順列を試して答えを見つける順列全探索。

この記事では、Python を使った順列全探索の基礎について説明します。その後、実際にAtCoderで出題された問題を順列全探索を使って解いてみます。

Python の文法については、次の記事にまとめています。

【競プロ】AtCoderを始めるためのPythonの基本

順列とは
順列全探索とは
1. itertools.permutations()
問題を解いてみる
まとめ
参考

順列とは

順列は、与えられた要素を並べ替えることで得られる組み合わせです。

要素の順序が異なれば、それぞれ異なる順列となります。

例えば、要素 A、B、C の順列は次の6つの組み合わせがあります。

順列全探索とは

順列全探索は、与えられた要素の集合から順列を生成し、それぞれの順列に対して処理を行う手法です。

Python では、itertools モジュールを使うことで、簡単に順列を生成できます。

import itertools

# 要素のリスト
A = ['A', 'B', 'C']

# 順列の生成
permutations = itertools.permutations(A)

# 生成した順列を１つずつ処理
for permutation in permutations:
    print(permutation)


"""実行結果
('A', 'B', 'C')
('A', 'C', 'B')
('B', 'A', 'C')
('B', 'C', 'A')
('C', 'A', 'B')
('C', 'B', 'A')
"""

itertools.permutations()

itertools.permutations() を使うことで、与えられた要素の順列を生成できます。

引数として渡す要素のリストの、すべての順列を生成してくれます。

生成された順列はイテレータとして返されるので、for で順番に処理することができます。

順列は、引数に応じた辞書順で出力されるので、引数のリストがソートされていれば、出力される順列もソートされた順番で生成されます。

次は、リストがソートされていないときの出力です。

import itertools

# 要素のリスト
A = ['C', 'A', 'B']

# 順列の生成
permutations = itertools.permutations(A)

# 順列を１つずつ処理
for permutation in permutations:
    print(permutation)


"""実行結果
('C', 'A', 'B')
('C', 'B', 'A')
('A', 'C', 'B')
('A', 'B', 'C')
('B', 'C', 'A')
('B', 'A', 'C')
"""

順列の長さを指定することもできます。r を指定していないときは、iterable の長さになります。

itertools.permutations(iterable, r=None)

先ほどのコードで r = 2 にしてみます。

A、B、C から２つ選んだ順列が生成されます。

import itertools

# 要素のリスト
A = ['C', 'A', 'B']

# 順列の生成
permutations = itertools.permutations(A, 2)

# 順列を１つずつ処理
for permutation in permutations:
    print(permutation)


"""実行結果
('C', 'A')
('C', 'B')
('A', 'C')
('A', 'B')
('B', 'C')
('B', 'A')
"""

問題を解いてみる

順列全探索を使って、ABC183 C – Travel を解いてみます。

問題概要

N 個の都市があります。都市 i から都市 j へ移動するには T_i,j の時間がかかります。

都市 1 を出発し、全ての都市をちょうど 1 度ずつ訪問してから都市 1 に戻るような経路のうち、移動時間の合計がちょうど K になるようなものはいくつありますか？

入力

N K

T_1,1 ・・・ T_1,N

：

T_N,1 ・・・ T_N,N

考え方

都市を訪れる順番について、順列全探索します。

N が４のとき、次のコードで要素「0, 1, 2, 3」の順列を生成できます。

都市１は「０」、都市２は「１」という風に都市の番号‐１で管理していることに注意してください。

import itertools

# 入力
N, K = map(int, input().split())
T = [list(map(int, input().split())) for _ in range(N)]

# 順列全探索
for permutation in itertools.permutations(range(N)):
    print(permutation)

"""実行結果（N が４のとき）
(0, 1, 2, 3)
(0, 1, 3, 2)
(0, 2, 1, 3)
(0, 2, 3, 1)
(0, 3, 1, 2)
(0, 3, 2, 1)
(1, 0, 2, 3)
(1, 0, 3, 2)
(1, 2, 0, 3)
(1, 2, 3, 0)
(1, 3, 0, 2)
(1, 3, 2, 0)
(2, 0, 1, 3)
(2, 0, 3, 1)
(2, 1, 0, 3)
(2, 1, 3, 0)
(2, 3, 0, 1)
(2, 3, 1, 0)
(3, 0, 1, 2)
(3, 0, 2, 1)
(3, 1, 0, 2)
(3, 1, 2, 0)
(3, 2, 0, 1)
(3, 2, 1, 0)
"""

都市１から出発なので、順列の１つ目が０以外は不要です。

次のように if 文で、不要な順列は飛ばします。

次のコードで、出発が都市１だけのときの順列が生成できます。

import itertools

# 入力
N, K = map(int, input().split())
T = [list(map(int, input().split())) for _ in range(N)]

# 順列全探索
for permutation in itertools.permutations(range(N)):

    # 出発が都市１以外のときはスキップ
    if permutation[0] != 0:
      continue
    
    print(permutation)

"""実行結果（N が４のとき）
(0, 1, 2, 3)
(0, 1, 3, 2)
(0, 2, 1, 3)
(0, 2, 3, 1)
(0, 3, 1, 2)
(0, 3, 2, 1)
"""

あとは、それぞれの順列についての移動時間を計算するだけです。

import itertools

# 入力
N, K = map(int, input().split())
T = [list(map(int, input().split())) for _ in range(N)]

answer = 0

# 順列全探索
for route in itertools.permutations(range(N)):
    
    # 出発が都市１以外のときはスキップ
    if route[0] != 0:
      continue
    
    # 移動時間
    time = 0

    # 経路の移動時間の計算
    for i in range(N-1):
      # 移動時間
      time += T[route[i]][route[i+1]]
    
    # 都市１に戻る移動時間
    time += T[route[N-1]][0]
    
    # 移動時間が K と等しいとき
    if time == K:
      answer += 1
      
# 答えの出力
print(answer)