クルトンのプログラミング教室

Pythonの使い方やPythonを使った競技プログラミングの解法などを解説しています。

Pythonで理解する蟻本「2-5 全点対最短路問題(ワーシャル-フロイド法)」(p.98)

この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「2-5 全点対最短路問題(ワーシャル-フロイド法)」(p.98)
のコードをPythonで書き直したものとなっています。

入力

蟻本には負の閉路がないグラフだけが入力例として載っていたので、負の閉路があるグラフも入力例として載せておきます。

また、蟻本に載っているグラフは無向グラフですが、このブログで追加した負の閉路があるグラフの入力例では有向グラフを使っています。

そのため、この記事では入力例は有向グラフに統一してあります。


V:頂点数\\E:辺数\\頂点f_iから頂点t_iに向かう重みc_iの有向辺が存在する(1\leq{i}\leq{E})


V\,E\\f_1\,t_1\,c_1\\:\\f_E\,t_E\,c_E

入力例1

f:id:kuruton456:20201031230032p:plain


7 20
0 1 2
1 0 2
0 2 5
2 0 5
1 2 4
2 1 4
1 3 6
3 1 6
1 4 10
4 1 10
2 3 2
3 2 2
3 5 1
5 3 1
4 5 3
5 4 3
4 6 5
6 4 5
5 6 9
6 5 9


入力例2

f:id:kuruton456:20201029231744p:plain


6 7
0 1 2
0 3 4
1 2 3
2 3 5
2 5 2
4 2 -4
5 4 1


コード

INF = float('inf')

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
d = [[INF] * V for _ in range(V)]    # d[u][v]は辺e=(u,v)のコスト(存在しない場合はINF、ただしd[i][i]=0とする)
for i in range(V):
    d[i][i] = 0
for _ in range(E):
    f, t, c = map(int,input().split())
    d[f][t] = c

def warshall_floyd():
    for k in range(V):
        for i in range(V):
            for j in range(V):
                d[i][j] = min(d[i][j], d[i][k] + d[k][j])

warshall_floyd()

print(d)

実行結果

上記のコードは実際に実行すると

[[0, 1, 2], [3, 4, 5], [6, 7, 8]]

のように一行で表示されます。

しかし、この形式はとても見にくいため下の実行例では

[[0, 1, 2],
 [3, 4, 5],
 [6, 7, 8]]

というように改行しています。

実際の実行結果とは異なりますが、リストの中身自体は同じなので問題ありません。

入力例1の場合
[[0,  2,  5,  7,  11, 8,  16],
 [2,  0,  4,  6,  10, 7,  15],
 [5,  4,  0,  2,  6,  3,  11],
 [7,  6,  2,  0,  4,  1,  9],
 [11, 10, 6,  4,  0,  3,  5],
 [8,  7,  3,  1,  3,  0,  8],
 [16, 15, 11, 9,  5,  8,  0]]
入力例2の場合
[[0,   2,   4,   4, 8,   6],
 [inf, 0,   2,   7, 6,   4],
 [inf, inf, -1,  4, 3,   1],
 [inf, inf, inf, 0, inf, inf],
 [inf, inf, -5,  0, -1,  -3],
 [inf, inf, -4,  1, 0,   -2]]


頂点iから頂点jに到達できない場合はd[i][j] = infとなり、負の閉路が存在する場合はd[i][i]が負となる頂点iが存在します。

これは上の実行例でも確かに成り立っています。