Data structures problem 30 of 30

Dijkstra's algorithm

Data structures
Hard

A weighted directed graph is stored as a dictionary of dictionaries: graph[a][b] is the cost of the edge from a to b. All costs are non-negative. Write dijkstra(graph, start) that returns the shortest distance from start to every reachable node.

graph = {
    "A": {"B": 4, "C": 1},
    "B": {"D": 1},
    "C": {"B": 2, "D": 5},
    "D": {},
}
dijkstra(graph, "A")
# {"A": 0, "C": 1, "B": 3, "D": 4}

Show hint

Always expand the unvisited node with the smallest known distance. A heap (heapq) gives you that node quickly.

Show solution

import heapq

def dijkstra(graph, start):
    dist = {start: 0}
    heap = [(0, start)]
    while heap:
        d, node = heapq.heappop(heap)
        if d > dist[node]:
            continue  # an outdated entry
        for nxt, cost in graph[node].items():
            nd = d + cost
            if nd < dist.get(nxt, float("inf")):
                dist[nxt] = nd
                heapq.heappush(heap, (nd, nxt))
    return dist

The direct edge A to B costs 4, but A to C to B costs 3, so B's distance is improved after C is processed. Rather than updating entries inside the heap, this version pushes a new entry and skips stale ones when they are popped. Runtime is O((V + E) log V). It needs non-negative costs; with negative edges, use Bellman-Ford. Map apps and network routing are built on this idea.

Want a problem like this every day? See today's problem and keep a streak going.