【ABC183 C】「Travel」を解く【Python3】

都市iから都市jへ移動するにはT_ijの時間がかかります。

N個の都市があるとき、都市1を出発して、すべての都市を一度ずつ訪問して都市1に戻るような経路のうち、移動時間の合計がKになるものがいくつあるか、求める問題です。

提出

import itertools

n, k = map(int, input().split())
t = [list(map(int,input().split())) for _ in range(n)]

li = list(range(2,n+1))
p_list = list(itertools.permutations(li))
ans = 0

for i in p_list:
    tra = 0
    tmp_li = i
    num = 1
    for j in tmp_li:
        tra += t[num-1][j-1]
        num = j
    tra += t[num-1][0]
    if tra == k:
        ans += 1

print(ans)

ありうる経路を全探索し、移動時間の合計がkになるものの個数を数えていきます。

最初と最後に訪問する都市1を除いた、都市2~Nまでの経路をすべて求めます。

「itertools.permutations()」を利用することで、順列のパターンを生成できます。

Nが4の場合、以下のようになります。

import itertools
n = 4
li = list(range(2, n+1))
p = list(itertools.permutations(li))
print(p)
# > [(2, 3, 4), (2, 4, 3), (3, 2, 4), (3, 4, 2), (4, 2, 3), (4, 3, 2)]

この場合、一つ目の(2, 3, 4)を使い、1→2→3→4→1 の移動時間を求めてkと比較することになります。

続けて、(2, 4, 3)、 (3, 2, 4)、…… と調べていきます。

提出したプログラム内では、リストtの中に、移動時間が入っているため、「t[num-1][j-1]」を「tra」に足していき、合計時間を求めています。

（最後に都市1に戻るため、「t[num-1][0]」を足しています）

合計時間が、kと等しい場合に「ans」に1を加えていきます。

最終的に「ans」に個数が入っているので、これを出力します。