ABC183 C問題 Travel をitertoolsで楽々と解く - 競技プログラミング悪戦苦闘日記

順列問題のABC183 C問題 Travelを解いて、解説します。

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

都市の数が $2 \leq N \leq 8$ なので、1回ずつ都市を巡回するパターンは最大で8!、せいぜい40320通りなので、全パターン列挙でいけると思いました。
実際には最初と最後が1であることが決まっているので、最大で7! = 5040通りになります。
あとはサンプルを見て順列列挙でいけると思いました。

まずは1を含んでいないルートを全パターン列挙します。

routes = list(permutations(range(2, N + 1), N - 1))

あとルートごとに前後に1を加えつつ、移動時間の合計を足します。
最後にKと一致していたら、カウントアップします。

from itertools import permutations

N, K = map(int, input().split())
t_map = [list(map(int, input().split())) for _ in range(N)]

# 最初と最後の1は含んでいない全ルートを列挙する
routes = list(permutations(range(2, N + 1), N - 1))

ok_routes = 0
for route in routes:
    # ルートの最初と最後に1を追加する
    r = [1]
    r.extend(list(route))
    r.extend([1])

    total = 0
    # ルートの合計時間を計算する
    for i in range(N):
        if total > K:
            break
        total += t_map[r[i] - 1][r[i + 1] - 1]
    # Kと一致してたらOK
    if total == K:
        ok_routes += 1
print(ok_routes)