65 lines
2.1 KiB
Python
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)
|