Initial commit

python/Algorithms/dijkstra_heap.py

@@ -0,0 +1,42 @@
+# 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()