43 lines
873 B
Python
43 lines
873 B
Python
# https://en.wikipedia.org/wiki/Dijkstra's_algorithm in O(E*log(V))
|
|
from heapq import *
|
|
|
|
|
|
class Edge:
|
|
def __init__(self, t, cost):
|
|
self.t = t
|
|
self.cost = cost
|
|
|
|
|
|
def dijkstra(graph, s):
|
|
n = len(graph)
|
|
pred = [-1] * n
|
|
prio = [float('inf')] * n
|
|
prio[s] = 0
|
|
q = [(0, s)]
|
|
while q:
|
|
(prio_u, u) = heappop(q)
|
|
if prio_u != prio[u]:
|
|
continue
|
|
for e in graph[u]:
|
|
v = e.t
|
|
nprio = prio[u] + e.cost
|
|
if prio[v] > nprio:
|
|
prio[v] = nprio
|
|
pred[v] = u
|
|
heappush(q, (nprio, v))
|
|
return prio, pred
|
|
|
|
|
|
def test():
|
|
g = [[] for _ in range(3)]
|
|
g[0].append(Edge(1, 3))
|
|
g[0].append(Edge(2, 2))
|
|
g[1].append(Edge(2, -2))
|
|
|
|
dist, pred = dijkstra(g, 0)
|
|
assert dist == [0, 3, 1]
|
|
assert pred == [-1, 0, 1]
|
|
|
|
|
|
test()
|