単一始点最短路探索アルゴリズム(1)

競技プログラミング

こんにちは、えぬわいです(普通)

ダイクストラ法の復習です。

import queue
push = queue.heappush
pop = queue.heappop
INF = 10000
from collections import defaultdict
G = defaultdict(list)
N, M, S, GOAL = map(int, input().split())
V = N
d = []
for i in range(M):
    f, t, c = map(int, input().split())
    G[f-1].append((t-1, c))
    G[t-1].append((f-1, c))

def dijkstra(s):
    global d
    q = []
    d = [INF] * V
    d[s] = 0
    push(q, (0, s))

    while q:
        cost, v = pop(q)
        print(f'd[{v}] = {cost}')
        if d[v] < cost:
            continue
        for to, tmpcost in G[v]:
            if d[to] > d[v] + tmpcost:
                d[to] = d[v] + tmpcost
                push(q, (d[to], to))

dijkstra(S-1)
print(d[GOAL-1])

蟻本を参考にして書くとこんな感じ。 queueを使えるところっていうのがよくわかった。他の自分で考えたものでも使えそうな予感。 テストケース

8 10 1 8
1 2 2
1 3 5
2 3 4
2 4 6
3 4 2
4 7 1
2 6 10
6 7 3
6 8 5
7 8 9
% python dijkstra.py < testcase.txt
d[0] = 0
d[1] = 2
d[2] = 5
d[3] = 7
d[3] = 8
d[6] = 8
d[5] = 11
d[5] = 12
d[7] = 16
d[7] = 17
16

おお〜。感慨深い。