Dijkstra’s Algorithm in Python: Shortest Path with heapq

BFS answers “fewest hops”. A maps app answers “cheapest trip”, and those are different questions the moment edges carry different costs. Take a four-node graph: A reaches B directly for 4, or goes through C for 1 + 2 = 3. BFS sees the direct edge, counts one hop, and happily reports the expensive route. Dijkstra’s algorithm fixes that with one idea and one data structure: always settle the closest unsettled node next, and use a heap to find it quickly.

Below is the version that is short enough to write on a whiteboard, a trace you can follow by hand, the line most people delete by accident, and the exact input that breaks the algorithm. All code was executed, and the final function was checked against Bellman-Ford on 2,000 random graphs.

The same graph, two answers

The graph used throughout this guide has edges A→B (4), A→C (1), C→B (2), C→D (7) and B→D (1):

g = {
    'A': [('B', 4), ('C', 1)],
    'C': [('B', 2), ('D', 7)],
    'B': [('D', 1)],
    'D': [],
}

BFS counts hops, so it reports B and C as one hop away and D as two. Dijkstra reports the costs. The real outputs side by side:

bfs hops: {'A': 0, 'B': 1, 'C': 1, 'D': 2}
dijkstra: {'A': 0, 'B': 3, 'C': 1, 'D': 4}

B costs 3, not 4, because the detour through C is cheaper. If you need a refresher on the BFS side, the traversal guide covers it.

The whole algorithm in 15 lines

import heapq

def dijkstra(graph, src):
    dist = {src: 0}
    heap = [(0, src)]
    while heap:
        d, u = heapq.heappop(heap)
        if d > dist.get(u, float('inf')):
            continue                      # stale entry, skip it
        for v, w in graph.get(u, []):
            nd = d + w
            if nd < dist.get(v, float('inf')):
                dist[v] = nd
                heapq.heappush(heap, (nd, v))
    return dist

Python’s heapq has no “decrease key” operation, so instead of updating an entry the code pushes a better one and leaves the old entry in the heap. That is called lazy deletion. The pushed tuple is (distance, node) because the heap orders by the first element. If you are new to the module, the heapq patterns guide explains heappush and heappop.

Trace it by hand

Follow the heap on the graph above. “Stale” means the popped distance is worse than the best known one.

Notice that B was pushed twice and D twice. Both outdated entries were harmless because the stale check threw them away.

The line people delete by accident

The if d > dist...: continue check looks optional, because the distances come out correct without it. They do, but the work does not. On a random graph with 300 nodes and 6,000 edges, counting how many times a node’s edges were scanned:

node expansions with the stale check:    300
node expansions without the stale check: 668
same distances either way: True

Without the check, a node is expanded again for every outdated entry. On dense graphs that quietly eats the complexity you are being asked to show. Keep the check.

Getting the path, not just the distance

Store the predecessor each time a distance improves, then walk backwards from the target. Stopping as soon as the target is popped is safe, because its distance is final at that moment.

def shortest_path(graph, src, dst):
    dist = {src: 0}
    prev = {}
    heap = [(0, src)]
    while heap:
        d, u = heapq.heappop(heap)
        if d > dist.get(u, float('inf')):
            continue
        if u == dst:
            break
        for v, w in graph.get(u, []):
            nd = d + w
            if nd < dist.get(v, float('inf')):
                dist[v] = nd
                prev[v] = u
                heapq.heappush(heap, (nd, v))
    if dst not in dist:
        return None
    path = [dst]
    while path[-1] != src:
        path.append(prev[path[-1]])
    return dist[dst], path[::-1]

shortest_path(g, 'A', 'D')   # -> (4, ['A', 'C', 'B', 'D'])

Two interview problems that are secretly Dijkstra

Network Delay Time

A signal starts at node k. How long until every node has it, or -1 if some node never does? Run Dijkstra from k and take the largest distance. If fewer than n nodes were reached, the answer is -1.

from collections import defaultdict

def network_delay_time(times, n, k):
    graph = defaultdict(list)
    for u, v, w in times:
        graph[u].append((v, w))
    dist = dijkstra(graph, k)
    return max(dist.values()) if len(dist) == n else -1

network_delay_time([[2,1,1],[2,3,1],[3,4,1]], 4, 2)   # -> 2
network_delay_time([[1,2,1]], 2, 2)                    # -> -1

Path With Minimum Effort

On a grid of heights, the effort of a path is its biggest single height jump, and you want the path that minimises it. Nothing says “graph” here, but each cell is a node and each step is an edge. The only change is the cost rule: instead of adding edge weights, take the maximum along the path. Dijkstra still works because that cost never decreases as the path grows.

def min_effort(heights):
    R, C = len(heights), len(heights[0])
    best = [[float('inf')] * C for _ in range(R)]
    best[0][0] = 0
    heap = [(0, 0, 0)]
    while heap:
        e, r, c = heapq.heappop(heap)
        if e > best[r][c]:
            continue
        if (r, c) == (R - 1, C - 1):
            return e
        for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            nr, nc = r + dr, c + dc
            if 0 <= nr < R and 0 <= nc < C:
                ne = max(e, abs(heights[nr][nc] - heights[r][c]))
                if ne < best[nr][nc]:
                    best[nr][nc] = ne
                    heapq.heappush(heap, (ne, nr, nc))

min_effort([[1,2,2],[3,8,2],[5,3,5]])   # -> 2
min_effort([[1,2,3],[3,8,4],[5,3,5]])   # -> 1

The input that breaks it: negative edges

Dijkstra’s correctness rests on one assumption: once a node is the closest unsettled one, no later path can beat it. A negative edge breaks that. Consider A→B (2), A→C (5) and C→B (−4). The true cost to B is 5 − 4 = 1. The common textbook version, which marks a node as done the first time it is popped, answers differently:

dijkstra with visited set: {'A': 0, 'B': 2, 'C': 5}
true shortest (Bellman-Ford): B = 1

B was finalised at 2 before the cheaper route through C was discovered. When a problem mentions negative weights, switch to Bellman-Ford, which relaxes every edge V − 1 times, or to a DAG-specific method using the order from topological sort.

Complexity in one breath

With a binary heap, every edge can cause at most one push, so the heap holds at most E entries and each push or pop costs O(log E), which is O(log V). Total time is O((V + E) log V), space is O(V + E). For the vocabulary around that, see Big O notation explained.

Questions interviewers ask after you finish

“Why does it need non-negative weights?”

Because the proof says the closest unsettled node is already final. A negative edge can later undercut it, as the example above shows.

“What if every edge weighs the same?”

Then it collapses into BFS, and the heap is wasted overhead. Use a plain queue. If weights are only 0 or 1, a deque (0-1 BFS) beats the heap.

“Could you stop earlier?”

Yes. For a single target, break when the target is popped. That is exactly what shortest_path does.

To keep practising, work through Network Delay Time, Path With Minimum Effort and Cheapest Flights Within K Stops, the last of which needs a stop counter in the state. Connectivity-only questions are cheaper with Union-Find, and the roadmap shows how graph algorithms fit into a full study plan.