Pythonで理解する蟻本「2-5 経路復元」(p.98) - クルトンのプログラミング教室

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 経路復元」(p.98)
のコードをPythonで書き直したものとなっています。

入力
入力例
隣接リストを用いたコード（O(|V|^2)）
- 実行結果

入力

ダイクストラ法は負の辺があると使えないため、負閉路の存在を考える必要はありません。

そのため、この記事では入力例は蟻本のテストケースのみにし、無向辺グラフとして入力を受け取っています。

$V:頂点数\\E:辺数\\頂点s_iと頂点t_iの間に重みc_iの無向辺が存在する(1\leq{i}\leq{E})$

$V\,E\\s_1\,t_1\,c_1\\:\\s_E\,t_E\,c_E$

入力例

f:id:kuruton456:20201029232302p:plain

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

隣接リストを用いたコード（O(|V|^2)）

INF = float('inf')

prev = []

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
cost = [[INF] * V for _ in range(V)]    # cost[u][v]は辺e=(u,v)のコスト(存在しない場所はINF)
for _ in range(E):
    s, t, c = map(int,input().split())
    cost[s][t] = c
    cost[t][s] = c


d = [INF for _ in range(V)]    # 最短距離
used = [False] * V    # すでに使われたかのフラグ
prev = [-1] * V    # 最短路の直前の頂点

# 始点sから各頂点への最短距離を求める
def dijkstra(s):
    d[s] = 0
    
    while True:
        v = -1
        for u in range(V):
            if not used[u] and (v == -1 or d[u] < d[v]):
                v = u
        
        if v == -1:
            break
        used[v] = True
        
        for u in range(V):
            if d[u] > d[v] + cost[v][u]:
                d[u] = d[v] + cost[v][u]
                prev[u] = v

# 頂点tへの最短路
def get_path(t):
    path = []
    while t != -1:
        path.append(t)    # tがsになるまでprev[t]を辿っていく
        t = prev[t]
    # このままだとt->sの順になっているので逆順にする
    path.reverse()
    return path

dijkstra(0)    # 頂点0を始点とする

path = get_path(6)    # 頂点0から頂点6への最短経路

print(path)

実行結果

[0, 2, 3, 5, 4, 6]

頂点0→2→3→5→4→6が最短経路となる。
f:id:kuruton456:20201105235135p:plain