Skip to content

Instantly share code, notes, and snippets.

@hanfang
Last active November 11, 2023 22:09
Show Gist options
  • Save hanfang/89d38425699484cd3da80ca086d78129 to your computer and use it in GitHub Desktop.
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
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"))
@JoelSebstn
Copy link

This is great,,, Thanks man

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment