ebisukeプログラミング初心者脱出黙示録

30歳を過ぎてから始めたプログラミングと競プロの記録。Pythonで取り組んでいます。Arduinoで電子工作も

AtCoder-ABC183 C - Travel【python解答例】

AtCoder Beginner Contest183 C問題に挑戦です。

実際にコンテストに出たときは解けませんでしたが、前回の記事のように順列の勉強を行いましたので再挑戦してみたいと思います。
ebisuke33.hatenablog.com

AtCoder Beginner Contest183 C - Travel

atcoder.jp

問題文

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

制約

・2≤N≤8
・i≠jのとき1≤Ti,j≤10^8
・Ti,i=0
・Ti,j=Tj,i
・1≤K≤10^9
・入力はすべて整数

解答例

import itertools

n, k = map(int,input().strip().split())
tlist = []

for i in range(n):
    array = list(map(int, input().strip().split()))
    tlist.append(array)

counter = 0

per_list = list(itertools.permutations(range(1,n))) # (※1)
# print(per_list)

for i in per_list:
    # print(i) (※2)
    i = [0] + list(i) + [0]
    # print(i) (※3)
    t_sum = 0
    for j in range(n):
        t_sum += tlist[i[j]][i[j+1]] # (※4)
    if t_sum == k:
        counter += 1

print(counter)

解説

順列を学びましたが簡単にはいきませんでしたので、細かい説明をしながら進めたいと思います。
まず、※1の部分では順列を作成しました。
ですがn = 4の場合、次の行で出力した結果は[(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]のようになります。
問題文から都市0から出発し、都市0に戻らなければならないのでその途中の経路を作成したイメージです。
私は公式解説だけではここが理解できず、かなり悩みました。
※2ではさきほど作成した順列のなかからひとつづつ取り出してループさせています。
例として(2,1,3)を取り出した場合で説明します。
※2の次の行における出力は(2,1,3)となり、都市2からスタートし、都市3に到着して終わりとなってしまいます。
したがって、i = [0] + list(i) + [0]の部分でリストの最初と最後に0を追加します。
これで※3の出力は(0,2,1,3,0)となり、都市0からスタートし、都市0に戻るという問題文通りの経路の順列を考えることができます。
※4ではt_sum変数にどんどん都市移動の時間を足していっています。
例に挙げた(0,2,1,3,0)では①tlist[0][2](都市0から2への移動時間)②tlist[2][1](都市2から1への移動時間)③tlist[1][3](都市1から3への移動時間)④tlist[3][0](都市3から0への移動時間)と順に足していきます。
その後、移動時間を足し終えたt_sumとkが一致していたらcounter変数を1増加させていき、すべての場合において確認が終わったら、counter変数を出力しています。
以上の考え方でプログラムを作るとACとなりました。

順列を学べば簡単にできるとたかをくくっていましたが、実際には問題にあわせて工夫が必要で難しかったです。
頭をつかって応用することが大事なんですね。

次はABC183のD問題に挑戦したいと思います。
学ぶことが多くて大変ですが、新しいことを覚えていく感覚を大事にして頑張りたいです。


ABC183についての以前の記事はこちら
ebisuke33.hatenablog.com