Dijkstra: Shortest Paths with heapq
30 min · Intermediate
Need help? Ask a tutor on WhatsApp (opens in a new tab)A delivery van needs the cheapest way from a depot to every town. Roads have different lengths, so breadth-first search is not enough. Dijkstra's algorithm finds all the cheapest distances, as long as no road has a negative cost.
The nearest unfinished node is already final
Dijkstra's algorithm finds the shortest distance from one start node to every other node in a graph whose edge weights are not negative, by always settling the closest unsettled node.
Keep a best-known distance for every node, starting at infinity except the start at 0. Repeatedly take the unfinished node with the smallest distance. No later route can beat it, because every road costs zero or more.
Example 1 · Simple version: scan for the nearest node
edges = [(0, 1, 4), (0, 2, 1), (2, 1, 2), (1, 3, 1), (2, 3, 5)]graph = [[] for _ in range(4)]for u, v, w in edges: graph[u].append((v, w)) graph[v].append((u, w))
INF = float('inf')dist = [INF] * 4dist[0] = 0done = [False] * 4for _ in range(4): u = min((i for i in range(4) if not done[i]), key=lambda i: dist[i]) done[u] = True for v, w in graph[u]: dist[v] = min(dist[v], dist[u] + w)print(dist)[0, 3, 1, 4]
What this code does
- Each round scans every node to find the closest unfinished one, so the time is O(V squared).
- It is fine for dense graphs and small V. For sparse graphs a heap is much faster.
Example 2 · Build the weighted adjacency list
edges = [(0, 1, 4), (0, 2, 1), (2, 1, 2), (1, 3, 1), (2, 3, 5)]graph = [[] for _ in range(4)]for u, v, w in edges: graph[u].append((v, w)) graph[v].append((u, w))
for u, row in enumerate(graph): print(u, row)0 [(1, 4), (2, 1)] 1 [(0, 4), (2, 2), (3, 1)] 2 [(0, 1), (1, 2), (3, 5)] 3 [(1, 1), (2, 5)]
What this code does
- An edge list becomes a list of
(neighbour, weight)pairs for each node. - A road works both ways, so each edge is added twice.
A priority queue gives the nearest node fast
Scanning all nodes for the smallest distance costs O(V) each time. A heap returns the smallest in O(log V). Push pairs (distance, node): Python compares tuples left to right, so the distance decides.
The code
When a shorter route to a node is found, push a new pair. The old pair stays in the heap, so skip it when it comes out with a larger distance than the best known. This is called lazy deletion.
def dijkstra(graph, start): dist = [float('inf')] * len(graph) dist[start] = 0 heap = [(0, start)] while heap: d, u = heapq.heappop(heap) if d > dist[u]: continue for v, w in graph[u]: if d + w < dist[v]: dist[v] = d + w heapq.heappush(heap, (d + w, v)) return distExample 3 · Dijkstra with heapq
import heapq
def dijkstra(graph, start): dist = [float('inf')] * len(graph) dist[start] = 0 heap = [(0, start)] while heap: d, u = heapq.heappop(heap) if d > dist[u]: continue for v, w in graph[u]: if d + w < dist[v]: dist[v] = d + w heapq.heappush(heap, (d + w, v)) return dist
edges = [(0, 1, 4), (0, 2, 1), (2, 1, 2), (1, 3, 1), (2, 3, 5)]graph = [[] for _ in range(4)]for u, v, w in edges: graph[u].append((v, w)) graph[v].append((u, w))
print(dijkstra(graph, 0))[0, 3, 1, 4]
What this code does
- Node 2 is reached first, at cost 1. From there node 1 costs 3, which beats the direct road of 4.
- Node 3 costs 4 through node 1, not 6 through node 2.
Example 4 · Rebuild the path
edges = [(0, 1, 4), (0, 2, 1), (2, 1, 2), (1, 3, 1), (2, 3, 5)]graph = [[] for _ in range(4)]for u, v, w in edges: graph[u].append((v, w)) graph[v].append((u, w))
import heapq
dist = [float('inf')] * 4parent = [-1] * 4dist[0] = 0heap = [(0, 0)]while heap: d, u = heapq.heappop(heap) if d > dist[u]: continue for v, w in graph[u]: if d + w < dist[v]: dist[v] = d + w parent[v] = u heapq.heappush(heap, (d + w, v))
path = []node = 3while node != -1: path.append(node) node = parent[node]print(dist[3], path[::-1])4 [0, 2, 1, 3]
What this code does
parent[v]remembers the node that gave the best distance. Walking back from the target and reversing the list gives the route.- Storing parents adds O(V) space and no extra time.
It needs import heapq. Here graph[u] holds (neighbour, weight) pairs. Each edge pushes at most once, so the time is O((V + E) log V) and the space is O(V + E).
Where it fails
The "already final" idea needs weights of zero or more. One negative road can make a longer-looking route cheaper later. Use Bellman-Ford then. For equal weights, plain breadth-first search is simpler and costs O(V + E).
Example 5 · Unreachable nodes
import heapq
def dijkstra(graph, start): dist = [float('inf')] * len(graph) dist[start] = 0 heap = [(0, start)] while heap: d, u = heapq.heappop(heap) if d > dist[u]: continue for v, w in graph[u]: if d + w < dist[v]: dist[v] = d + w heapq.heappush(heap, (d + w, v)) return dist
graph = [[(1, 2)], [(0, 2)], [(3, 1)], [(2, 1)]]dist = dijkstra(graph, 0)print(dist)print([d if d != float('inf') else -1 for d in dist])[0, 2, inf, inf] [0, 2, -1, -1]
What this code does
- Nodes 2 and 3 form a separate island, so their distance stays infinity.
- Many problems print -1 for "no route". Test for
float('inf')before printing.
Example 6 · Why negative weights break it
edges = {0: [(1, 4), (2, 5)], 1: [], 2: [(1, -3)]}dist = {0: 0, 1: 4, 2: 5}done = {0}for u in [1, 2]: done.add(u) for v, w in edges[u]: if v not in done: dist[v] = min(dist[v], dist[u] + w)print(dist){0: 0, 1: 4, 2: 5}What this code does
- Node 1 was settled at 4. The route 0 → 2 → 1 costs 5 - 3 = 2, but the settled value is never updated.
- The true distance to node 1 is 2, so a settled node is not final when weights can be negative.
Next step
Check your understanding
11 quick questions