Graphs: Shortest Paths, Union-Find and Topological Sort

The previous chapter covered traversals. This chapter covers the three graph tools that solve the next tier of problems: Dijkstra's algorithm for weighted shortest paths, union-find (disjoint sets) for dynamic connectivity, and topological sort for dependencies in directed acyclic graphs. Each solves a specific kind of question, and recognising the question is most of the work.

1. Which tool for which question

2. Dijkstra's algorithm

Dijkstra finds the shortest distance from a source to every node when edge weights are non-negative. It is BFS with a priority queue: always expand the unfinished node with the smallest known distance. Once a node is popped with its smallest distance, that distance is final (a path through any other unvisited node would be at least as long, because weights are non-negative).

import heapq
from collections import defaultdict

def dijkstra(edges, start, n):
    graph = defaultdict(list)
    for a, b, w in edges:
        graph[a].append((b, w))
    dist = {i: float("inf") for i in range(n)}
    dist[start] = 0
    heap = [(0, start)]
    while heap:
        d, node = heapq.heappop(heap)
        if d > dist[node]:                  # a stale entry: a shorter path was already found
            continue
        for nxt, w in graph[node]:
            nd = d + w
            if nd < dist[nxt]:
                dist[nxt] = nd
                heapq.heappush(heap, (nd, nxt))
    return dist

edges = [(0, 1, 4), (0, 2, 1), (2, 1, 2), (1, 3, 1), (2, 3, 5)]
assert dijkstra(edges, 0, 4) == {0: 0, 1: 3, 2: 1, 3: 4}

Time $O((V+E)\log V)$ with a binary heap. Space $O(V+E)$. The check d > dist[node] skips outdated heap entries, which is how Python's heap handles the lack of a decrease-key operation.

Network delay time. A signal is sent from node $k$. How long until all nodes receive it? Run Dijkstra and take the maximum distance, or -1 if some node is unreachable.

def network_delay_time(times, n, k):
    graph = defaultdict(list)
    for u, v, w in times:
        graph[u].append((v, w))
    dist = {}
    heap = [(0, k)]
    while heap:
        d, node = heapq.heappop(heap)
        if node in dist:
            continue
        dist[node] = d                          # first pop of a node is its shortest distance
        for nxt, w in graph[node]:
            if nxt not in dist:
                heapq.heappush(heap, (d + w, nxt))
    return max(dist.values()) if len(dist) == n else -1

assert network_delay_time([[2, 1, 1], [2, 3, 1], [3, 4, 1]], 4, 2) == 2
assert network_delay_time([[1, 2, 1]], 2, 1) == 1
assert network_delay_time([[1, 2, 1]], 2, 2) == -1

Dijkstra applies to more than road maps. Path with minimum effort (the cost of a path is the largest height difference along it) and cheapest flights within $k$ stops use the same idea with a different notion of "cost" or an extra state dimension. Anything with positive costs and "find the cheapest route" is a candidate.

Negative weights. Dijkstra can be wrong with negative edges. Bellman-Ford relaxes every edge $V-1$ times and handles negatives and detects negative cycles, in $O(VE)$.

def bellman_ford(n, edges, source):
    dist = [float("inf")] * n
    dist[source] = 0
    for _ in range(n - 1):
        for u, v, w in edges:
            if dist[u] + w < dist[v]:
                dist[v] = dist[u] + w
    for u, v, w in edges:                      # one more pass: any improvement means a negative cycle
        if dist[u] + w < dist[v]:
            return None
    return dist

assert bellman_ford(4, [(0, 1, 4), (0, 2, 5), (1, 2, -3), (2, 3, 2)], 0) == [0, 4, 1, 3]
assert bellman_ford(3, [(0, 1, 1), (1, 2, -1), (2, 0, -1)], 0) is None

Floyd-Warshall computes all-pairs shortest paths in $O(V^3)$ by trying each node as an intermediate: it is three nested loops and is appropriate only when $V$ is a few hundred at most.

def floyd_warshall(n, edges):
    d = [[float("inf")] * n for _ in range(n)]
    for i in range(n):
        d[i][i] = 0
    for u, v, w in edges:
        d[u][v] = min(d[u][v], w)
    for k in range(n):                          # intermediate node
        for i in range(n):
            for j in range(n):
                if d[i][k] + d[k][j] < d[i][j]:
                    d[i][j] = d[i][k] + d[k][j]
    return d

dm = floyd_warshall(3, [(0, 1, 5), (1, 2, 3), (0, 2, 10)])
assert dm[0][2] == 8

3. Union-find (disjoint set union)

Union-find maintains a collection of disjoint sets and supports two operations: find(x) returns the representative of $x$'s set, and union(a, b) merges two sets. With two optimisations, path compression and union by size or rank, each operation is effectively constant time (the inverse Ackermann function, which is below 5 for any practical input).

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n
        self.components = n

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]     # path compression (halving)
            x = self.parent[x]
        return x

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False                                       # already connected
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra                                   # attach the smaller tree under the larger
        self.size[ra] += self.size[rb]
        self.components -= 1
        return True

uf = UnionFind(5)
assert uf.union(0, 1) and uf.union(1, 2)
assert uf.find(0) == uf.find(2) and uf.find(3) != uf.find(0)
assert uf.union(0, 2) is False
assert uf.components == 3

When union-find beats DFS: the graph is built incrementally (edges arrive over time) and you need connectivity queries along the way. DFS would be rerun after every change, while union-find answers in near-constant time.

Redundant connection. A tree with one extra edge: find the edge that creates a cycle. Process edges in order, and the first edge whose endpoints are already connected is the answer.

def find_redundant_connection(edges):
    uf = UnionFind(len(edges) + 1)
    for a, b in edges:
        if not uf.union(a, b):
            return [a, b]

assert find_redundant_connection([[1, 2], [1, 3], [2, 3]]) == [2, 3]
assert find_redundant_connection([[1, 2], [2, 3], [3, 4], [1, 4], [1, 5]]) == [1, 4]

Number of provinces, accounts merge, longest consecutive sequence, satisfiability of equality equations all reduce to grouping items. For accounts merge, union emails that appear in the same account, then group by representative.

Minimum spanning tree (Kruskal). Sort edges by weight and add each edge unless it creates a cycle (checked with union-find). The total weight of the chosen edges is the minimum cost to connect all nodes.

def kruskal(n, edges):
    uf = UnionFind(n)
    total, used = 0, 0
    for u, v, w in sorted(edges, key=lambda e: e[2]):
        if uf.union(u, v):
            total += w
            used += 1
    return total if used == n - 1 else -1             # -1 if the graph is not connected

assert kruskal(4, [(0, 1, 1), (1, 2, 2), (2, 3, 3), (0, 3, 10), (0, 2, 4)]) == 6
assert kruskal(3, [(0, 1, 1)]) == -1

Time $O(E\log E)$ dominated by sorting.

4. Topological sort

A topological ordering of a directed graph lists every node before the nodes it points to. It exists only for a DAG. It models dependencies: course prerequisites, build order, task scheduling. If the graph has a cycle, no ordering exists, which is how you detect deadlocks or impossible schedules.

Kahn's algorithm (BFS)

Track each node's in-degree (number of unmet prerequisites). Start with nodes that have in-degree 0. Repeatedly remove one, appending it to the order, and decrement the in-degree of its neighbours. If at the end the order contains all nodes, it is valid. If not, a cycle blocked some nodes.

from collections import deque

def topological_order(n, prerequisites):
    """prerequisites: (course, requires) pairs. Returns an order, or [] if there is a cycle."""
    graph = defaultdict(list)
    indegree = [0] * n
    for course, req in prerequisites:
        graph[req].append(course)
        indegree[course] += 1
    queue = deque(i for i in range(n) if indegree[i] == 0)
    order = []
    while queue:
        node = queue.popleft()
        order.append(node)
        for nxt in graph[node]:
            indegree[nxt] -= 1
            if indegree[nxt] == 0:
                queue.append(nxt)
    return order if len(order) == n else []

def can_finish(n, prerequisites):
    return len(topological_order(n, prerequisites)) == n

assert can_finish(2, [(1, 0)]) is True
assert can_finish(2, [(1, 0), (0, 1)]) is False
order = topological_order(4, [(1, 0), (2, 0), (3, 1), (3, 2)])
assert order.index(0) < order.index(1) < order.index(3) and order.index(0) < order.index(2) < order.index(3)

Time $O(V+E)$.

DFS with three states

Alternatively, run DFS and add each node to the order after its descendants (postorder), then reverse. Detect cycles by tracking nodes currently on the recursion path: reaching one again means a back edge, hence a cycle.

def topological_order_dfs(n, prerequisites):
    graph = defaultdict(list)
    for course, req in prerequisites:
        graph[req].append(course)
    state = [0] * n                      # 0 = unvisited, 1 = on the current path, 2 = done
    order = []
    def visit(node):
        if state[node] == 1:
            return False                 # a cycle
        if state[node] == 2:
            return True
        state[node] = 1
        for nxt in graph[node]:
            if not visit(nxt):
                return False
        state[node] = 2
        order.append(node)
        return True
    for i in range(n):
        if not visit(i):
            return []
    return order[::-1]

assert topological_order_dfs(2, [(0, 1), (1, 0)]) == []
o = topological_order_dfs(3, [(1, 0), (2, 1)])
assert o == [0, 1, 2]

Alien dictionary

You are given words sorted in an unknown alphabet. Derive letter order from adjacent word pairs: the first position where two adjacent words differ gives one ordering edge. Then topologically sort the letters.

def alien_order(words):
    graph = {c: set() for w in words for c in w}
    for a, b in zip(words, words[1:]):
        if len(a) > len(b) and a.startswith(b):
            return ""                                    # invalid: a longer word before its prefix
        for x, y in zip(a, b):
            if x != y:
                graph[x].add(y)
                break
    indegree = {c: 0 for c in graph}
    for c in graph:
        for nxt in graph[c]:
            indegree[nxt] += 1
    queue = deque(c for c in graph if indegree[c] == 0)
    out = []
    while queue:
        c = queue.popleft()
        out.append(c)
        for nxt in graph[c]:
            indegree[nxt] -= 1
            if indegree[nxt] == 0:
                queue.append(nxt)
    return "".join(out) if len(out) == len(graph) else ""

assert alien_order(["wrt", "wrf", "er", "ett", "rftt"]) == "wertf"
assert alien_order(["z", "x"]) == "zx"
assert alien_order(["abc", "ab"]) == ""

5. Putting it together: choosing the tool

Cheapest flights within at most $k$ stops. Weighted, with a constraint on the number of edges. Dijkstra's greedy finalisation does not respect the stops limit, so use Bellman-Ford limited to $k+1$ rounds, or BFS by layers.

def cheapest_flights(n, flights, src, dst, k):
    prices = [float("inf")] * n
    prices[src] = 0
    for _ in range(k + 1):                      # at most k+1 edges
        nxt = prices[:]                         # read old values, write new ones
        for u, v, w in flights:
            if prices[u] + w < nxt[v]:
                nxt[v] = prices[u] + w
        prices = nxt
    return -1 if prices[dst] == float("inf") else prices[dst]

flights = [[0, 1, 100], [1, 2, 100], [2, 0, 100], [1, 3, 600], [2, 3, 200]]
assert cheapest_flights(4, flights, 0, 3, 1) == 700
assert cheapest_flights(4, flights, 0, 3, 2) == 400
assert cheapest_flights(3, [[0, 1, 100], [1, 2, 100]], 0, 2, 0) == -1

The lesson: the constraint in the problem decides the algorithm. Unrestricted positive weights suggest Dijkstra. A limit on edges suggests rounds of Bellman-Ford or a state of (node, steps).

6. Common mistakes

• Using Dijkstra with negative weights. It can return wrong answers. State the assumption.
• Skipping the stale-entry check in Dijkstra, which can cause repeated work or wrong distances.
• Using BFS on a weighted graph (it counts edges, not cost).
• Forgetting path compression and union by size in union-find, degrading it to a linear chain.
• Reversing the edge direction in topological sort (prerequisite pairs are easy to flip). Draw one example.
• Not detecting cycles in topological sort. Compare the order length with $n$.
• Using visited marks only for directed cycle detection. You need the "on the current path" state.
• Counting components wrong in union-find: decrement only when a union actually merges two different sets.

7. Practice set

1. Network delay time, cheapest flights within $k$ stops, path with minimum effort.
2. Path with maximum probability, swim in rising water (Dijkstra or union-find).
3. Number of provinces, redundant connection, accounts merge.
4. Number of connected components, graph valid tree, most stones removed.
5. Course schedule I and II, minimum height trees, parallel courses.
6. Alien dictionary, sequence reconstruction.
7. Min cost to connect all points (Prim or Kruskal).
8. Evaluate division (weighted union-find or graph DFS).
9. Find the city with the smallest number of neighbours at a threshold (Floyd-Warshall).
10. Critical connections in a network (bridges, a harder DFS idea using discovery times).