Graph algorithms in Python

Dijkstra's Algorithm Lab

Shortest paths from one vertex to every other, with a min heap doing the hard part. Step through the code and watch the priority queue, the distance table and each edge relaxation change together.

The problem: single-source shortest paths

graph = {
    'A': [('B', 4), ('C', 1)],
    'B': [('D', 1)],
    'C': [('B', 2), ('D', 5), ('E', 8)],
    'D': [('E', 3), ('F', 6)],
    'E': [('F', 1)],
    'F': [],
}
dist, prev = dijkstra(graph, 'A')

Each vertex lists (neighbour, weight) pairs. From one source, find the cheapest total weight to every vertex (dist) and the vertex just before it on that route (prev). Weights must be non-negative: distances, times, costs.

The greedy idea

# repeat until nothing is left:
#   u = the unfinished vertex with the smallest dist
#   u is now final
#   try to improve u's neighbours through u

The unfinished vertex with the smallest tentative distance can't be improved later: any other route to it has to leave through a vertex that is already at least as far away, and with weights ≥ 0 a detour can only add. So it is safe to lock it in.

Relaxation: try a shortcut

nd = dist[u] + w      # route to v through u
if nd < dist[v]:      # better than the best so far?
    dist[v] = nd
    prev[v] = u       # remember how we got here

When u becomes final, each edge u → v asks one question: is going through u cheaper than the best route to v found so far? If yes, dist[v] drops and prev[v] points back at u. That is the whole update step.

Why a min heap

pq = []
heapq.heappush(pq, (4, 'B'))
heapq.heappush(pq, (1, 'C'))
heapq.heappop(pq)     # (1, 'C'): smallest d first

"Smallest unfinished dist" is a min query, and heapq is a min heap (see the Heap Lab): push and pop are O(log n). Tuples compare by their first item, so (d, v) pops by distance. heapq can't lower an entry already inside, so the code pushes a new one and skips the old one later.

not reached (dist ∞)reached, waiting in pqfinal (popped)u: being processedsourceedge being relaxed (dashed: no improvement)prev edges: the shortest-path tree~~(4, 'B')~~stale pq entry

Source

Graph number under each vertex = dist

dist and prev

Tip: the ← and → keys step too.

Why it costs O((V + E) log V)

Count the work by what each line can do in total, not per loop turn:

Each vertex is processed once. The stale check lets only one entry per vertex through, so for v, w in graph[u] runs once per vertex and checks every edge exactly once overall: E checks.
Each edge can push at most once. A push only happens on a successful relaxation, so pq receives at most E + 1 entries (the +1 is the source).
Every entry is popped at most once, stale or not, so there are at most E + 1 pops too.
Each push or pop costs O(log(heap size)), and the heap never holds more than E + 1 entries. Since E ≤ V², log E ≤ 2 log V, which is O(log V).

Add it up: O(V) to set up dist and prev, O(E) edge checks, and O(E) heap operations at O(log V) each: O((V + E) log V), often written O(E log V) when every vertex is reachable. Space is O(V + E): the two dicts plus up to E + 1 heap entries.

Counts measured by the lab
Run	V	E	Pushes	Pops	Stale pops	Edge checks
graph from A (Run Dijkstra)	6	9	10	10	4	9
small from S (Stale entries)	4	5	6	6	2	5
neg from A (Negative edge)	3	3	4	4	0	3

On the main graph every one of the 9 edges improved something, so pq got the maximum E + 1 = 10 pushes, and 4 of the 10 pops were stale. A heap with a decrease-key operation (a Fibonacci heap) avoids the duplicates and reaches O(E + V log V), but Python's heapq has no decrease-key, and lazy deletion is simpler and fast in practice.

Dijkstra next to BFS and Bellman-Ford

Algorithm	Edge weights	Frontier	Time	Why it works
BFS	none (every edge costs 1)	collections.deque, first in first out	O(V + E)	Vertices come out in order of hop count, so the first time you reach a vertex is the shortest. Dijkstra with all weights equal does the same thing, just slower.
Dijkstra	≥ 0	heapq min heap of (d, v)	O((V + E) log V)	Vertices come out in order of distance. A popped vertex is final because detours can only add weight.
Bellman-Ford	any, including negative	none: relax every edge, V − 1 rounds	O(V · E)	After round k every shortest path of k edges is correct. One extra round that still improves something means a negative cycle.

Common mistakes

Marking a vertex visited when it is pushed

visited = {source}
while pq:
    d, u = heapq.heappop(pq)
    for v, w in graph[u]:
        if v not in visited:
            visited.add(v)       # locked in at the FIRST route
            dist[v] = d + w
            prev[v] = u
            heapq.heappush(pq, (dist[v], v))

done = set()
while pq:
    d, u = heapq.heappop(pq)
    if u in done:                # same job as d > dist[u]
        continue
    done.add(u)                  # final only when popped
    for v, w in graph[u]:
        nd = d + w
        if nd < dist[v]:
            ...

This is BFS thinking: the first route found to a vertex is not the cheapest one. On this lab's graph from A, B gets locked at 4 by the direct edge before C finds 1 + 2 = 3, and the error spreads: it returns B = 4, D = 6, E = 9, F = 12 instead of B = 3, D = 4, E = 7, F = 8. A vertex is only final when it is popped.

Leaving out the stale check

while pq:
    d, u = heapq.heappop(pq)
    # no stale check here
    for v, w in graph[u]:        # re-runs for every old entry
        ...

    if d > dist[u]:
        continue

With non-negative weights the answer is still right: an old entry has d > dist[u], so d + w can't beat what u's earlier, smaller pop already set. But every stale pop re-scans all of u's edges (4 wasted edge checks on the main graph, on top of 9), and on a graph with many improvements into a high-degree vertex that work multiplies. A done set checked at pop time does the same job.

A plain list and `min()`

pq = [(0, source)]
while pq:
    d, u = min(pq)               # O(len(pq)) scan
    pq.remove((d, u))            # and another one
    ...

    d, u = heapq.heappop(pq)     # O(log V)

Correct, but every pop scans the whole list instead of costing O(log V). The classic array version (scan all V vertices for the smallest unfinished dist) is O(V²): fine for a dense graph where E is close to V², far slower than the heap on the sparse graphs you usually meet.

Negative weights

neg = {'A': [('B', 4), ('C', 2)],
       'B': [('C', -3)],
       'C': []}
dijkstra(neg, 'A')   # pops C at 2 before B is seen

def bellman_ford(graph, source):
    dist = {v: float('inf') for v in graph}
    dist[source] = 0
    for _ in range(len(graph) - 1):
        for u in graph:
            for v, w in graph[u]:
                if dist[u] + w < dist[v]:
                    dist[v] = dist[u] + w
    return dist

The greedy step finalises C at 2, but A → B → C costs 4 − 3 = 1. The lazy code on this page happens to recover by popping C twice (step through Negative edge above), but then a vertex can be processed many times, the time bound is gone, and a negative cycle never ends. Use Bellman-Ford when any weight is negative.

Cost of each piece

Piece	Time	Why
Set up dist and prev	O(V)	one entry per vertex
heappush / heappop	O(log V)	sift up / down a heap of at most E + 1 entries
Stale check	O(1)	one comparison, then continue
All relaxations	O(E)	each edge checked once, when its tail is popped
Whole run	O((V + E) log V)	at most E + 1 pushes and pops
Rebuild one path from prev	O(V)	one step per vertex on the path
Space	O(V + E)	dist, prev and the heap

Only need one target? Return as soon as it is popped (if u == target: return d): its distance is final at that moment. Same worst case, often much less work.