algo/python/44_shortest_path/shortest_path.py
2019-01-06 21:03:03 +00:00

65 lines
2.1 KiB
Python

"""
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)