ABC183 C-Travel解説

解説しようと思ったきっかけ

競プロを初めてまだ間もないですが、競プロをやってきて初めて(標準)ライブラリを使ったため、新鮮味があって楽しかったから。使わなくても当然解けますが...

問題

ABC183 Travel

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

入力例
4 330
0 1 10 100
1 0 20 200
10 20 0 300
100 200 300 0
出力例
2

import itertools

#入力を受け取る
N, K = map(int, input().split())
T = [ [None] for i in range(N) ]
for i in range(N):
    T[i] = list(map(int, input().split()))

発想

スタートは都市1固定で、最後も都市1に帰ってこなければならないので、
Nまでの順列を考え、全組み合わせの中で最後の要素が1である組み合わせのみ考えれば良い。
例としてN=3のときの並べ方の組み合わせを見てみる。

for i in itertools.permutations(range(1, 3+1)):
    print(i)

出力
(1, 2, 3)
(1, 3, 2)
(2, 1, 3)
(2, 3, 1) <-最後に都市1へ帰る組
(3, 1, 2)
(3, 2, 1) <-最後に都市1へ帰る組

最初に都市1が来たらまずいと思うが、はじめは都市1スタート固定かつ最後の要素が都市1で条件分岐しているので、問題ない。よって、コードは以下のようになる。

res = 0

for i in itertools.permutations(range(1, N+1)):
    i = list(i) #itertoolはタプル型なのでリストに変換
    if i[-1] == 1: #最後の要素が1の時(最後に都市1に帰ってくる)のみ通過
        tmp = T[0][i[0]-1] #スタートは都市1固定なので初めだけ処理
        for j in range(1, N):
            tmp += T[i[j-1]-1][i[j]-1]
        if tmp == K: #合計が目的のKならカウント
            res += 1

まとめ

二重目のfor文の中身はややこしいのでノートに実際に書いたり、デバックを走らせればわかると思います。自分もノートに書いてやってました!