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

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

Pythonで理解する蟻本「2-5 単一始点最短路問題2(ダイクストラ法)」(p.96)

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

入力

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

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


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')

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]