この記事は拙著「AtCoder 凡人が『緑』になるための精選50問詳細解説」のサンプルです
価格：100円
kindle：https://www.amazon.co.jp/dp/B09C3TPQYV
booth(pdf)：https://sano192.booth.pm/items/3179185
前：https://qiita.com/sano192/items/b9967d40093c4c3be42d
次：https://qiita.com/sano192/items/d6d2397f5f1bde156631
【目標】
・itertoolsを使えるようになる
【概要】
競技プログラミングにおいて「全てのパターンを試す」というのは基本であるが、そのためのコードを書くことは意外と難しい。itertoolsを使えば順列、組み合わせ、重複組合せなどさまざまなパターンを簡単に作ることができる。使えると便利どころか使えないと解けない問題も多いので、この問題でマスターしよう。
【方針】
制約が2<=N<=8とかなり小さいため、Nが最大のケースで全ての経路を考えても
(2~8の順列の数)=7!=5040通り
これなら全ての経路について移動時間の合計がKとなったものをカウントするという方針で十分間に合う。
問題は全ての経路をどのように列挙するか。
そこで活躍するのがitertools。
順列を列挙してくれるライブラリで、これを使えばこの問題はいとも簡単に解ける。
順列の列挙は以下のように書く。
# 1からNまでの順列
import itertools
for root in itertools.permutations(range(1,N+1)):
こう書くとrootに順列を小さい方から順に全通り格納してくれる。
と、言われてもよくわからないと思う。
試しにN=4として、rootを出力してみよう。
N=4
# 1からNまでの順列
import itertools
for root in itertools.permutations(range(1,N+1)):
print(root)
「出力結果」
(1, 2, 3, 4)
(1, 2, 4, 3)
(1, 3, 2, 4)
(1, 3, 4, 2)
(1, 4, 2, 3)
(1, 4, 3, 2)
(2, 1, 3, 4)
(2, 1, 4, 3)
(2, 3, 1, 4)
(2, 3, 4, 1)
(2, 4, 1, 3)
(2, 4, 3, 1)
(3, 1, 2, 4)
(3, 1, 4, 2)
(3, 2, 1, 4)
(3, 2, 4, 1)
(3, 4, 1, 2)
(3, 4, 2, 1)
(4, 1, 2, 3)
(4, 1, 3, 2)
(4, 2, 1, 3)
(4, 2, 3, 1)
(4, 3, 1, 2)
(4, 3, 2, 1)
1,2,3,4を使って作れる24個の順列が列挙された。
この順列を経路と対応させることで全経路を探索できる。
作った順列と対応する経路の移動時間を計算し、Kとなるものを数えて出力すればよい。
【実装】
入力を受け取る。
N,K=map(int, input().split())
次に各都市間の移動時間=Tだが、これは二次元配列として受け取る。
まずtimeという空のリストを作る。
次に一行ずつリストで受け取り、それをtimeに格納していく。
time=[]
for i in range(N):
T=list(map(int, input().split()))
time.append(T)
このように受け取ることで
都市A→都市Bの移動時間=time[A-1][B-1]
と取り出す事ができる。
(0インデックスで受け取りしているため、都市1は0、都市2は1、...都市AはA-1、都市BはB-1、...というように都市番号を-1する必要がある)
例えば都市5→都市7の移動時間は
time[4][6]
で参照できる。
次にansという変数を作る。これは移動時間の合計=Kとなった経路の数をカウントする変数。最初は0。
ans=0
ここから経路を順に試していく。itertoolsの出番だ。
まずitertoolsをインポートする。
import itertools
都市1がスタート地点だが受け取りが0インデックスのため、-1して0からスタート。
そして都市2~都市N=1~(N-1)を1度ずつ通って最後に0へ戻る。
つまり
0→(1~N-1の順列)→0
というルートをすべて試す。
ここからの説明は少々ややこしいので説明の後に「具体例」を書いている。わからなかったらそちらも見てほしい。
必要なのは(1~N-1の順列)だからまず以下のように書く。
for root in itertools.permutations(range(1,N)):
(1)0→root[0]
now_time=0
now_time+=time[0][root[0]]
now_timeは今いる場所に至るまでにかかった時間。
0→root[0]に向かう場合、移動時間はtime[0][root[0]]で参照できる。
now_timeに移動時間(time[0][root[0]])を足す。
now_city=root[0]
now_cityは今いる都市の番号。
0→root[0]と移動してroot[0]にたどり着いたからnow_city=root[0]に更新。
(2)root[0]→root[1]、root[1]→root[2]、root[2]→root[3]、...、root[N-3]→root[N-2]
作られた順列の順に都市を移動していく。
for i in range(1,N-1):
to_city=root[i]
now_time+=time[now_city][to_city]
now_city=to_city
to_cityは次に行く都市の番号。
now_city→to_cityの移動時間を参照してnow_timeに足す。
移動したのでnow_city=to_cityと更新。
(3)root[N-2]→0
最後に都市1へ戻る=root[N-2]→0へ移動する。
移動時間をnow_timeに足す。
now_time+=time[now_city][0]
これで
0→(1~N-1の順列)→0
の移動時間が計算できた。
「具体例」
N=4
root=(2,3,1)
のパターンを考える。
つまり
root[0]=2
root[1]=3
root[2]=1
これは
都市1→(root)→都市1つまり
都市1→都市3→都市4→都市2→都市1
のパターン。
(0インデックスだから、rootの中身は「都市番号-1」となっていることに注意しよう)
(1)0→root[0]
now_time=0
now_time+=time[0][root[0]]
まずスタート地点都市1→都市3。
つまり0→root[0]=0→2への移動。
now_time+=time[0][root[0]] (=time[0][2])
で0→2への移動時間をnow_timeに足す。
(2)root[0]→root[1]、root[1]→root[2]、root[2]→root[3]、...、root[N-3]→root[N-2]
for i in range(1,N-1):
to_city=root[i]
now_time+=time[now_city][to_city]
now_city=to_city
i=1
to_city=root[1] (=3)
(行き先は都市4)
now_time+=time[now_city][to_city] (=time[2][3])
(今いる場所：都市3→行き先：都市4 の移動時間をnow_timeに足す)
now_city=to_city (=3)
(移動したから今いる場所を都市4へ更新)
i=2
to_city=root[2] (=1)
(行き先は都市2)
now_time+=time[now_city][to_city] (=time[3][1])
(今いる場所：都市4→行き先：都市2 の移動時間をnow_timeに足す)
now_city=to_city(=1)
(移動したから今いる場所を都市2へ更新)
(3)root[N-2]→0
now_time+=time[now_city][0]
最後に都市1へ戻る。
now_city→0、つまり1→0への移動。
(都市2→都市1への移動)
now_time+=time[now_city][0] (=time[1][0])
で1→0への移動時間をnow_timeに足す。
「具体例終わり」
移動時間=now_timeがKならばansにカウントする。
now_time+=time[now_city][0]
if now_time==K:
ans+=1
これを全ての順列に対して行い、最後に答えを出力して終わり。
print(ans)
itertoolsは
・順列
・重複なしの組み合わせ
・重複ありの組み合わせ
・直積
が作れる。
それぞれの書き方を以下に記載するのでコンテスト中にコピペできるよう保存しておこう。
import itertools
# 順列
#(1,2,3),(1,3,2),(2,1,3),(2,3,1),...,(3,2,1)
for seq in itertools.permutations(range(1,4)):
# 重複なしの組み合わせ
# (1,2,3),(1,2,4),...,(7,8,9)
for seq in itertools.combinations(range(1,10),3):
# 重複ありの組み合わせ
#(1,1,1),(1,1,2),...,(9,9,9)
for seq in itertools.combinations_with_replacement(range(1,10),3):
# 直積
#(1,1),(1,2),(1,3),(2,1),(2,2),...,(3,3)
for seq in itertools.product(range(1,4),range(1,4)):
【コード全文】
N,K=map(int, input().split())
time=[]
for i in range(N):
T=list(map(int, input().split()))
time.append(T)
ans=0
import itertools
for root in itertools.permutations(range(1,N)):
now_time=0
now_time+=time[0][root[0]]
now_city=root[0]
for i in range(1,N-1):
to_city=root[i]
now_time+=time[now_city][to_city]
now_city=to_city
now_time+=time[now_city][0]
if now_time==K:
ans+=1
print(ans)
この記事は拙著「AtCoder 凡人が『緑』になるための精選50問詳細解説」のサンプルです
価格：100円
kindle：https://www.amazon.co.jp/dp/B09C3TPQYV
booth(pdf)：https://sano192.booth.pm/items/3179185
前：https://qiita.com/sano192/items/b9967d40093c4c3be42d
次：https://qiita.com/sano192/items/d6d2397f5f1bde156631