Dijkstra shortest path done, using heapq.

Pretty simple after I first implemented in Java.

1 files changed, 45 insertions, 0 deletions

diff --git a/sources/dijkstra.py b/sources/dijkstra.py
new file mode 100644
index 0000000..e5160d4
--- /dev/null
+++ b/sources/dijkstra.py
@@ -0,0 +1,45 @@
+# Uses python3
+
+import sys
+import heapq
+
+
+def distance(adj, cost, s, t):
+    if s == t:
+        return 0
+    dist = [float('inf')] * len(adj)
+    visited = [False] * len(adj)
+    dist[s] = 0
+    q = []
+    heapq.heappush(q, (dist[s], s))
+
+    while len(q) != 0:
+        u = heapq.heappop(q)[1]
+        if visited[u]:
+            continue
+        visited[u] = True
+        if len(adj[u]) > 0:
+            for i in range(len(adj[u])):
+                if dist[adj[u][i]] > dist[u] + cost[u][i]:
+                    dist[adj[u][i]] = dist[u] + cost[u][i]
+                    heapq.heappush(q, (dist[adj[u][i]], adj[u][i]))
+
+    if dist[t] == float('inf'):
+        return -1
+    return dist[t]
+
+
+if __name__ == '__main__':
+    input = sys.stdin.read()
+    data = list(map(int, input.split()))
+    n, m = data[0:2]
+    data = data[2:]
+    edges = list(zip(zip(data[0:(3 * m):3], data[1:(3 * m):3]), data[2:(3 * m):3]))
+    data = data[3 * m:]
+    adj = [[] for _ in range(n)]
+    cost = [[] for _ in range(n)]
+    for ((a, b), w) in edges:
+        adj[a - 1].append(b - 1)
+        cost[a - 1].append(w)
+    s, t = data[0] - 1, data[1] - 1
+    print(distance(adj, cost, s, t))