← Advanced Graphs

Drill #95 — Dijkstra simplified

Medium BFS + Heap Advanced Graphs

In plain English: Find the shortest path from one starting node to all other nodes in a graph with positive edge weights.

Always process the closest unfinished node first (via a min-heap). Once a node is popped, its shortest distance is finalized — greedy on distance.

Prompt

Given a weighted directed graph as an adjacency list and a source node, find the shortest distance to all other nodes using Dijkstra's algorithm.

Try to write it from scratch before scrolling down.

Solution

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
        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

# graph: {0: [(1,4),(2,1)], 1: [(3,1)], 2: [(1,2),(3,5)], 3: []}
# dijkstra(graph, 0) => {0:0, 1:3, 2:1, 3:4}
O((V+E) log V) time · O(V) space

Related Micro Drills

← Drill #94 Drill #96 →