Pythonで理解する蟻本「2-5 単一始点最短路問題2(ダイクストラ法)」(p.96)
この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「2-5 単一始点最短路問題2(ダイクストラ法)」(p.96)
のコードをPythonで書き直したものとなっています。
入力例
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')
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 # すでに使われたかのフラグ
# 始点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):
d[u] = min(d[u], d[v] + cost[v][u])
dijkstra(0) # 頂点0を始点とする
print(d)
実行結果
[0, 2, 5, 7, 11, 8, 16]
優先度付きキューを使った解法(O(|E| log |V|))
# 優先度付きキューをインポートする
import heapq
INF = float('inf')
V, E = map(int,input().split()) # Vは頂点数、Eは辺数
G = [[] for _ in range(V)]
for _ in range(E):
s, t, c = map(int,input().split())
G[s].append([t, c])
G[t].append([s, c])
d = [INF for _ in range(V)]
def dijkstra(s):
d[s] = 0
que = []
heapq.heappush(que, (0, s))
while len(que) != 0:
p = heapq.heappop(que)
v = p[1]
if d[v] < p[0]:
continue
for i in range(len(G[v])):
e = G[v][i]
if d[e[0]] > d[v] + e[1]:
d[e[0]] = d[v] + e[1]
heapq.heappush(que, (d[e[0]], e[0]))
dijkstra(0) # 頂点0を始点とする
print(d)
実行結果
[0, 2, 5, 7, 11, 8, 16]