Merge pull request #223 from jerryderry/shortest-path-python

shortest path algorithm in python

python/44_shortest_path/shortest_path.py

@@ -0,0 +1,64 @@
+"""
+    Dijkstra algorithm
+
+    Author: Wenru Dong
+"""
+
+from dataclasses import dataclass
+from queue import PriorityQueue
+
+@dataclass
+class Edge:
+    start_id: int
+    end_id: int
+    weight: int
+
+@dataclass(order=True)
+class Vertex:
+    distance_to_start = float("inf")
+    vertex_id: int
+
+class Graph:
+    def __init__(self, num_vertices: int):
+        self._num_vertices = num_vertices
+        self._adjacency = [[] for _ in range(num_vertices)]
+    
+    def add_edge(self, from_vertex: int, to_vertex: int, weight: int) -> None:
+        self._adjacency[from_vertex].append(Edge(from_vertex, to_vertex, weight))
+
+    def dijkstra(self, from_vertex: int, to_vertex: int) -> None:
+        vertices = [Vertex(i) for i in range(self._num_vertices)]
+        vertices[from_vertex].distance_to_start = 0
+        visited = [False] * self._num_vertices
+        predecessor = [-1] * self._num_vertices
+        q = PriorityQueue()
+        q.put(vertices[from_vertex])
+        visited[from_vertex] = True
+        while not q.empty():
+            min_vertex = q.get()
+            if min_vertex.vertex_id == to_vertex:
+                break
+            for edge in self._adjacency[min_vertex.vertex_id]:
+                next_vertex = vertices[edge.end_id]
+                if min_vertex.distance_to_start + edge.weight < next_vertex.distance_to_start:
+                    next_vertex.distance_to_start = min_vertex.distance_to_start + edge.weight
+                    predecessor[next_vertex.vertex_id] = min_vertex.vertex_id
+                    if not visited[next_vertex.vertex_id]:
+                        q.put(next_vertex)
+                        visited[next_vertex.vertex_id] = True
+            
+        path = lambda x: path(predecessor[x]) + [str(x)] if from_vertex != x else [str(from_vertex)]
+        print("->".join(path(to_vertex)))
+
+
+if __name__ == "__main__":
+    graph = Graph(6)
+    graph.add_edge(0, 1, 10)
+    graph.add_edge(0, 4, 15)
+    graph.add_edge(1, 2, 15)
+    graph.add_edge(1, 3, 2)
+    graph.add_edge(2, 5, 5)
+    graph.add_edge(3, 2, 1)
+    graph.add_edge(3, 5, 12)
+    graph.add_edge(4, 5, 10)
+    graph.dijkstra(0, 5)