-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
112 lines (87 loc) · 3.16 KB
/
main.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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import queue
from collections import namedtuple
'''
Dijkstra's algorithm not only calculates the shortest (lowest weight) path on a graph
from source vertex S to destination V, but also calculates the shortest path from S
to every other vertex.
This implementation in Python doesn't return the shortest paths to all vertices,
but it could. It returns the shortest path from source to destination,
and the sum of the weights along that path.
https://googleyasheck.com/dijstras-algorithm-in-python-3/
John Washam https://googleyasheck.com/
'''
Edge = namedtuple('Edge', ['vertex', 'weight'])
class GraphUndirectedWeighted(object):
def __init__(self, vertex_count):
self.vertex_count = vertex_count
self.adjacency_list = [[] for _ in range(vertex_count)]
def add_edge(self, source, dest, weight):
assert source < self.vertex_count
assert dest < self.vertex_count
self.adjacency_list[source].append(Edge(dest, weight))
self.adjacency_list[dest].append(Edge(source, weight))
def get_edge(self, vertex):
for e in self.adjacency_list[vertex]:
yield e
def get_vertex(self):
for v in range(self.vertex_count):
yield v
def dijkstra(graph, source, dest):
q = queue.PriorityQueue()
parents = []
distances = []
start_weight = float("inf")
for i in graph.get_vertex():
weight = start_weight
if source == i:
weight = 0
distances.append(weight)
parents.append(None)
q.put(([0, source]))
while not q.empty():
v_tuple = q.get()
v = v_tuple[1]
for e in graph.get_edge(v):
candidate_distance = distances[v] + e.weight
if distances[e.vertex] > candidate_distance:
distances[e.vertex] = candidate_distance
parents[e.vertex] = v
# primitive but effective negative cycle detection
if candidate_distance < -1000:
raise Exception("Negative cycle detected")
q.put(([distances[e.vertex], e.vertex]))
shortest_path = []
end = dest
while end is not None:
shortest_path.append(end)
end = parents[end]
shortest_path.reverse()
return shortest_path, distances[dest]
def main():
g = GraphUndirectedWeighted(9)
g.add_edge(0, 1, 4)
g.add_edge(1, 7, 6)
g.add_edge(1, 2, 1)
g.add_edge(2, 3, 3)
g.add_edge(3, 7, 1)
g.add_edge(3, 4, 2)
g.add_edge(3, 5, 1)
g.add_edge(4, 5, 1)
g.add_edge(5, 6, 1)
g.add_edge(6, 7, 2)
g.add_edge(6, 8, 2)
g.add_edge(7, 8, 2)
# for testing negative cycles
# g.add_edge(1, 9, -5)
# g.add_edge(9, 7, -4)
shortest_path, distance = dijkstra(g, 0, 1)
assert shortest_path == [0, 1] and distance == 4
shortest_path, distance = dijkstra(g, 0, 8)
assert shortest_path == [0, 1, 2, 3, 7, 8] and distance == 11
shortest_path, distance = dijkstra(g, 5, 0)
assert shortest_path == [5, 3, 2, 1, 0] and distance == 9
shortest_path, distance = dijkstra(g, 1, 1)
assert shortest_path == [1] and distance == 0
print(dijkstra(g, 0, 3))
if __name__ == "__main__":
main()