概要
前回の記事において、単一始点最短経路(Single Source Shortest Path; SSSP)の重みを求めるプログラムである Dijkstra(ダイクストラ)法のpython
による実装例を紹介しました。そこでは、隣接リストを用いて実装をしましたが、隣接行列を用いても同様に実装することができるので紹介しておきます。
実装
from heapq import heappush, heappop
def dijkstra(s, n): # (始点, ノード数)
dist = [INF] * n
hq = [(0, s)] # (distance, node)
dist[s] = 0
seen = [False] * n
while hq:
v = heappop(hq)[1] # node
seen[v] = True
for to in range(n):
if W[v][to] == INF: continue # 隣接する辺がない場合は処理をスキップ
if seen[to] == False and dist[v] + W[v][to] < dist[to]:
dist[to] = dist[v] + W[v][to]
heappush(hq, (dist[to], to))
return dist
実行イメージ
次のようなグラフを例に考えてみましょう。
- 入力例:
4 6 1 # ノード数, エッジ数, 始点ノード
0 1 1
0 2 4
2 0 1
1 2 2
3 1 1
3 2 5
- 入力の受け取り・隣接行列
W
の構築:
v, e, r = map(int, input().split()) # ノード数, エッジ数, 始点
INF = 10 ** 4 - 1 # 入力に応じて変更する
W = [[INF] * v for _ in range(v)] # 重み付き隣接行列
for i in range(e):
s, t, d = map(int, input().split())
W[s][t] = d
- 出力:
>>> print(*W, sep='\n') # 隣接行列
[999, 1, 4, 999]
[999, 999, 2, 999]
[1, 999, 999, 999]
[999, 1, 5, 999]
>>> d = dijkstra(r, v)
>>> print(d)
[3, 0, 2, 999]