Refactor Dijkstra pseudo-code and function for readability and type s…

graphs/dijkstra.py

@@ -3,23 +3,28 @@
 
 DIJKSTRA(graph G, start vertex s, destination vertex d):
 
-//all nodes initially unexplored
-
-1 -  let H = min heap data structure, initialized with 0 and s [here 0 indicates
-     the distance from start vertex s]
-2 -  while H is non-empty:
-3 -    remove the first node and cost of H, call it U and cost
-4 -    if U has been previously explored:
-5 -      go to the while loop, line 2 //Once a node is explored there is no need
-         to make it again
-6 -    mark U as explored
-7 -    if U is d:
-8 -      return cost // total cost from start to destination vertex
-9 -    for each edge(U, V): c=cost of edge(U,V) // for V in graph[U]
-10 -     if V explored:
-11 -       go to next V in line 9
-12 -     total_cost = cost + c
-13 -     add (total_cost,V) to H
+// all nodes initially unexplored
+
+1 -  let H = min heap data structure, initialized with (0, s)
+     // 0 is the distance from start vertex s
+2 -  let costs = dictionary to store minimum costs to reach each node,
+     initialized with {s: 0}
+3 -  while H is non-empty:
+4 -    remove the first node and cost from H, call them U and cost
+5 -    if U has been previously explored:
+6 -      continue // skip further processing and go back to while loop, line 3
+7 -    mark U as explored
+8 -    if U is d:
+9 -      return cost // total cost from start to destination vertex
+10 -   for each neighbor V and edge cost c of U in G:
+11 -     if V has been previously explored:
+12 -       continue to next neighbor V in line 10
+13 -     total_cost = cost + c
+14 -     if total_cost is less than costs.get(V, ∞):
+15 -       update costs[V] to total_cost
+16 -       add (total_cost, V) to H
+
+// At the end, if destination d is not reachable, return -1
 
 You can think at cost as a distance where Dijkstra finds the shortest distance
 between vertices s and v in a graph G. The use of a min heap as H guarantees
@@ -34,7 +39,7 @@
 import heapq
 
 
-def dijkstra(graph, start, end):
+def dijkstra(graph: dict[str, list[tuple[str, int]]], start: str, end: str) -> int:
     """Return the cost of the shortest path between vertices start and end.
 
     >>> dijkstra(G, "E", "C")
@@ -44,31 +49,37 @@ def dijkstra(graph, start, end):
     >>> dijkstra(G3, "E", "F")
     3
     """
+    heap: list[tuple[int, str]] = [(0, start)]  # (cost, node)
+    visited: set[str] = set()
+    costs: dict[str, int] = {start: 0}  # Store minimum costs to reach each node
 
-    heap = [(0, start)]  # cost from start node,end node
-    visited = set()
     while heap:
-        (cost, u) = heapq.heappop(heap)
+        cost, u = heapq.heappop(heap)
         if u in visited:
             continue
         visited.add(u)
         if u == end:
             return cost
+
         for v, c in graph[u]:
             if v in visited:
                 continue
-            next_item = cost + c
-            heapq.heappush(heap, (next_item, v))
+            next_cost = cost + c
+            # Only push to heap if a cheaper path is found
+            if next_cost < costs.get(v, float("inf")):
+                costs[v] = next_cost
+                heapq.heappush(heap, (next_cost, v))
+
     return -1
 
 
 G = {
-    "A": [["B", 2], ["C", 5]],
-    "B": [["A", 2], ["D", 3], ["E", 1], ["F", 1]],
-    "C": [["A", 5], ["F", 3]],
-    "D": [["B", 3]],
-    "E": [["B", 4], ["F", 3]],
-    "F": [["C", 3], ["E", 3]],
+    "A": [("B", 2), ("C", 5)],
+    "B": [("A", 2), ("D", 3), ("E", 1), ("F", 1)],
+    "C": [("A", 5), ("F", 3)],
+    "D": [("B", 3)],
+    "E": [("B", 4), ("F", 3)],
+    "F": [("C", 3), ("E", 3)],
 }
 
 r"""
@@ -80,10 +91,10 @@ def dijkstra(graph, start, end):
     ----------------- 3 --------------------
 """
 G2 = {
-    "B": [["C", 1]],
-    "C": [["D", 1]],
-    "D": [["F", 1]],
-    "E": [["B", 1], ["F", 3]],
+    "B": [("C", 1)],
+    "C": [("D", 1)],
+    "D": [("F", 1)],
+    "E": [("B", 1), ("F", 3)],
     "F": [],
 }
 
@@ -96,12 +107,12 @@ def dijkstra(graph, start, end):
     -------- 2 ---------> G ------- 1 ------
 """
 G3 = {
-    "B": [["C", 1]],
-    "C": [["D", 1]],
-    "D": [["F", 1]],
-    "E": [["B", 1], ["G", 2]],
+    "B": [("C", 1)],
+    "C": [("D", 1)],
+    "D": [("F", 1)],
+    "E": [("B", 1), ("G", 2)],
     "F": [],
-    "G": [["F", 1]],
+    "G": [("F", 1)],
 }
 
 short_distance = dijkstra(G, "E", "C")