-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdijkstra.py
More file actions
44 lines (36 loc) · 1 KB
/
dijkstra.py
File metadata and controls
44 lines (36 loc) · 1 KB
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
import heapq
graph = {
"A": {"B": 5, "C": 1},
"B": {"A": 5, "C": 2, "D": 1},
"C": {"A": 1, "B": 2, "D": 4, "E": 8},
"D": {"B": 1, "C": 4, "E": 3, "F": 6},
"E": {"C": 8, "D": 3},
"F": {"D": 6}
}
def init_dist(graph, s):
distance = {}
for i in graph:
if i == s:
distance[i] = 0
else:
distance[i] = float('inf')
return distance
def dijkstra(graph, s):
distance = init_dist(graph, s)
seen = set()
pq = []
heapq.heappush(pq, (0, s))
parent = {}
while pq:
dist, vertex = heapq.heappop(pq)
seen.add(vertex)
for w in graph[vertex]:
if w not in seen:
if dist + graph[vertex][w] < distance[w]:
distance[w] = dist + graph[vertex][w]
parent[w] = vertex
heapq.heappush(pq, (distance[w], w))
return parent, distance
if __name__ == "__main__":
parent, distance = dijkstra(graph, "A")
print(parent, distance)