-
Notifications
You must be signed in to change notification settings - Fork 0
/
dijsktra.py
82 lines (72 loc) · 2.72 KB
/
dijsktra.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
from collections import defaultdict
from typing import List, Dict, Set
class Graph():
def __init__(self) -> None:
"""
self.edges is a dict of all possible next nodes
e.g. {'X': ['A', 'B', 'C', 'E'], ...}
self.weights has all the weights between two nodes,
with the two nodes as a tuple as the key
e.g. {('X', 'A'): 7, ('X', 'B'): 2, ...}
"""
self.edges: Dict[str, List[str]] = defaultdict(list)
self.weights: Dict[Tuple[str,str], int] = {}
def add_edge(self, from_node: str, to_node: str, weight: int) -> None:
# Note: assumes edges are bi-directional
self.edges[from_node].append(to_node)
self.edges[to_node].append(from_node)
self.weights[(from_node, to_node)] = weight
self.weights[(to_node, from_node)] = weight
def dijsktra(graph: Graph, initial: str, end: str):
# shortest paths is a dict of nodes, whose value is a tuple of (previous node, weight)
shortest_paths: Dict[str, Tuple[str,int]] = {initial: (None, 0)}
current_node: str = initial
visited: Set[str] = set()
while current_node != end:
visited.add(current_node)
destinations = graph.edges[current_node]
weight_to_current_node = shortest_paths[current_node][1]
for next_node in destinations:
weight = graph.weights[(current_node, next_node)] + weight_to_current_node
if next_node not in shortest_paths:
shortest_paths[next_node] = (current_node, weight)
else:
current_shortest_weight = shortest_paths[next_node][1]
if current_shortest_weight > weight:
shortest_paths[next_node] = (current_node, weight)
next_destinations = {node: shortest_paths[node] for node in shortest_paths if node not in visited}
if not next_destinations:
return "Route Not Possible"
# next node is the destination with the lowest weight
current_node = min(next_destinations, key=lambda k: next_destinations[k][1])
# Work back through destinations in shortest path
path = []
while current_node is not None:
path.append(current_node)
current_node = shortest_paths[current_node][0]
# Reverse path
return path[::-1]
edges = [
('X', 'A', 7),
('X', 'B', 2),
('X', 'C', 3),
('X', 'E', 4),
('A', 'B', 3),
('A', 'D', 4),
('B', 'D', 4),
('B', 'H', 5),
('C', 'L', 2),
('D', 'F', 1),
('F', 'H', 3),
('G', 'H', 2),
('G', 'Y', 2),
('I', 'J', 6),
('I', 'K', 4),
('I', 'L', 4),
('J', 'L', 1),
('K', 'Y', 5),
]
graph: Graph = Graph()
for edge in edges:
graph.add_edge(*edge)
print(dijsktra(graph, 'X', 'Y'))