Union-Find and minimum spanning trees

Implement disjoint sets with path compression and union by rank, then use them for Kruskal, cycle detection and connectivity queries.

Disjoint set union

class DSU:
    """Union-Find with path compression and union by size.

    Both optimisations together give amortised O(alpha(n)) per operation,
    where alpha is the inverse Ackermann function: at most 4 for any
    input that fits in memory, so effectively constant.
    """
    __slots__ = ("parent", "size", "components")

    def __init__(self, n: int) -> None:
        self.parent = list(range(n))
        self.size = [1] * n
        self.components = n

    def find(self, x: int) -> int:
        root = x
        while self.parent[root] != root:
            root = self.parent[root]
        while self.parent[x] != root:      # path compression, iterative
            self.parent[x], x = root, self.parent[x]
        return root

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

    def connected(self, a: int, b: int) -> bool:
        return self.find(a) == self.find(b)


def count_components(n: int, edges: list[tuple[int, int]]) -> int:
    dsu = DSU(n)
    for a, b in edges:
        dsu.union(a, b)
    return dsu.components


def has_cycle(n: int, edges: list[tuple[int, int]]) -> bool:
    """Union-Find detects a cycle in an undirected graph edge by edge."""
    dsu = DSU(n)
    for a, b in edges:
        if not dsu.union(a, b):
            return True
    return False


if __name__ == "__main__":
    print(count_components(5, [(0, 1), (1, 2), (3, 4)]))   # 2
    print(has_cycle(4, [(0, 1), (1, 2), (2, 0)]))          # True

• Path compression and rollback are incompatible: compression rewrites many parents, which makes an undo log expensive. Use union by size alone if you need to undo.
• find must be called on the same structure for both arguments of a union, and the result must be recomputed after a union. Caching a root across a union is a classic bug.
• Union-Find answers connectivity and cycle detection but cannot answer path queries. It has no notion of distance, unlike BFS or Dijkstra.
• Path compression changes the parent array while iterating, so do not iterate over parent and call find at the same time.

Kruskal and Prim

import heapq
from typing import Sequence

Edge = tuple[int, int, int]     # weight, u, v

def kruskal(n: int, edges: Sequence[Edge]) -> tuple[int, list[Edge]]:
    """Sort by weight, add an edge whenever it joins two components."""
    ordered = sorted(edges)
    dsu = DSU(n)
    total = 0
    chosen: list[Edge] = []

    for w, u, v in ordered:
        if dsu.union(u, v):                # union returns False if already connected
            total += w
            chosen.append((w, u, v))
            if len(chosen) == n - 1:       # a spanning tree is complete here
                break
    return total, chosen


def prim(n: int, adj: dict[int, list[tuple[int, int]]], start: int = 0) -> int:
    """Grow one tree: repeatedly take the cheapest edge leaving it."""
    seen = [False] * n
    pq: list[tuple[int, int]] = [(0, start)]
    total = 0
    used = 0

    while pq and used < n:
        w, u = heapq.heappop(pq)
        if seen[u]:
            continue                       # a stale entry: the node is already in
        seen[u] = True
        total += w
        used += 1
        for v, weight in adj.get(u, ()):
            if not seen[v]:
                heapq.heappush(pq, (weight, v))
    return total if used == n else -1       # -1 when the graph is disconnected


if __name__ == "__main__":
    edges = [(1, 0, 1), (4, 0, 2), (2, 1, 2), (3, 1, 3), (5, 2, 3)]
    print(kruskal(4, edges))                # (6, [(1,0,1), (2,1,2), (3,1,3)])

    adj = {0: [(1, 1), (2, 4)], 1: [(0, 1), (2, 2), (3, 3)], 2: [(0, 4), (1, 2), (3, 5)], 3: [(1, 3), (2, 5)]}
    print(prim(4, adj))                     # 6

Both algorithms are greedy and both are correct, which surprises people the first time: the cut property guarantees that the cheapest edge crossing any partition belongs to some minimum spanning tree, and each algorithm only ever chooses such an edge.

Where minimum spanning trees do not apply

def minimum_spanning_forest(n: int, edges: Sequence[Edge]) -> tuple[int, int]:
    """Kruskal on a disconnected graph: a tree per component.

    Returns the total weight and the number of trees. Typical use:
    cluster points by connecting everything closer than a threshold.
    """
    ordered = sorted(edges)
    dsu = DSU(n)
    total = 0
    for w, u, v in ordered:
        if dsu.union(u, v):
            total += w
    return total, dsu.components


def single_linkage_clusters(n: int, edges: Sequence[Edge], k: int) -> int:
    """Stop Kruskal early to get exactly k clusters: the usual clustering trick."""
    ordered = sorted(edges)
    dsu = DSU(n)
    for w, u, v in ordered:
        if dsu.components == k:
            break
        dsu.union(u, v)
    return dsu.components


# What an MST does NOT give you:
#   - shortest paths. A minimum spanning tree minimises total edge weight,
#     not the distance between any two vertices in particular.
#   - a solution for directed graphs. That is the arborescence problem,
#     solved by Chu-Liu/Edmonds, not by Kruskal or Prim.
#   - a Steiner tree. Allowing extra junction vertices changes the problem
#     and makes it NP-hard.
#   - robustness to a single edge failing. A minimum bottleneck spanning
#     tree or a 2-edge-connected structure is the right answer there.

FAQ

Related

Graph algorithms: BFS, DFS and shortest paths Greedy algorithms and intervals

Last refreshed 2026-09-18.