Last active
December 4, 2024 10:11
-
-
Save hanfang/89d38425699484cd3da80ca086d78129 to your computer and use it in GitHub Desktop.
Finding the shortest path in a weighted DAG with Dijkstra in Python and heapq
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import collections | |
import heapq | |
def shortestPath(edges, source, sink): | |
# create a weighted DAG - {node:[(cost,neighbour), ...]} | |
graph = collections.defaultdict(list) | |
for l, r, c in edges: | |
graph[l].append((c,r)) | |
# create a priority queue and hash set to store visited nodes | |
queue, visited = [(0, source, [])], set() | |
heapq.heapify(queue) | |
# traverse graph with BFS | |
while queue: | |
(cost, node, path) = heapq.heappop(queue) | |
# visit the node if it was not visited before | |
if node not in visited: | |
visited.add(node) | |
path = path + [node] | |
# hit the sink | |
if node == sink: | |
return (cost, path) | |
# visit neighbours | |
for c, neighbour in graph[node]: | |
if neighbour not in visited: | |
heapq.heappush(queue, (cost+c, neighbour, path)) | |
return float("inf") | |
if __name__ == "__main__": | |
edges = [ | |
("A", "B", 7), | |
("A", "D", 5), | |
("B", "C", 8), | |
("B", "D", 9), | |
("B", "E", 7), | |
("C", "E", 5), | |
("D", "E", 15), | |
("D", "F", 6), | |
("E", "F", 8), | |
("E", "G", 9), | |
("F", "G", 11) | |
] | |
print ("Find the shortest path with Dijkstra") | |
print (edges) | |
print ("A -> E:") | |
print (shortestPath(edges, "A", "E")) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
This is great,,, Thanks man