Lesson 40 · Hard mode

Union-Find — connectivity over streams

📖 Walk me through it — plain English

Union-Find (also called DSU, short for "Disjoint Set Union") answers one question really fast: "are these two items in the same group?" — and lets you merge two groups together. "Disjoint" just means the groups never overlap; every item belongs to exactly one group. We care about this because lots of problems boil down to tracking connectivity as connections arrive one at a time: are two computers on the same network, are two people in the same friend circle, would adding this edge create a cycle?

The trick: each group is stored as a tree, and the tree's root is the group's "name." The array p (parent) holds, for each item, who its parent is. A root is an item that is its own parent (p[x] == x). To find which group an item is in, you walk up the parent links until you hit the root. Two items are in the same group exactly when they reach the same root.

Analogy: think of company org charts. Everyone reports to a boss, who reports to a bigger boss, up to the CEO at the top. To check if two employees are in the same company, follow each one's chain of bosses up to the CEO and see if it's the same CEO. Merging two companies = point one CEO at the other. Two extra tricks keep the charts short: union by rank hangs the shorter tree under the taller one so chains don't grow tall, and path compression re-points nodes closer to the root every time you walk up, flattening the chart over time.

Let's trace it on 5 items, 0..4. Each box shows an item's parent. A box outlined in green is a root (it points to itself). We'll run a few unions, then a query.

Start: every item is its own group, so every parent points to itself. All 5 are roots. (item index above, parent value in the box)

Step 1 · union(0, 1). Roots 0 and 1 have equal rank, so hang 1 under 0 and bump 0's rank to 1. Now p[1] = 0, so item 1 (highlighted) points at root 0.

Step 2 · union(2, 3). Same story on the other side: hang 3 under 2, bump 2's rank to 1. Now p[3] = 2. We have two groups of two: {0,1} rooted at 0, and {2,3} rooted at 2.

Step 3 · union(0, 2). find(0)=0, find(2)=2 — different roots, so merge. Both roots have rank 1 (a tie), so we keep 0 as the parent, set p[2] = 0, and bump 0's rank to 2. Item 2 (highlighted) now points at root 0. Note item 3 still points at 2, not directly at 0 — yet.

Step 4 · find(3) — and watch path compression. We start at 3: p[3]=2 (not itself), so we set p[3] = p[p[3]] = p[2] = 0, then move to 0; p[0]=0, stop. The walk returns root 0, AND it re-pointed item 3 straight at 0. The chart got flatter for free.

Step 5 · union(0, 3) now. find(0)=0 and find(3)=0 — same root, so they're already connected. The code returns False and changes nothing. (That False is exactly how you detect a cycle: the edge you're adding would link two things already joined.)

Why it's so fast: union by rank keeps trees from ever getting tall, and path compression flattens them more every time you call find. Together they make both find and union run in effectively constant time — formally O(α(n)), where α (the inverse Ackermann function) grows so unbelievably slowly that for any input you'll ever see in practice it's at most 4. The handy mental model: it's basically O(1) per operation, so processing a stream of n connections is about O(n) overall. That's why this template is worth memorizing: it turns "are these connected?" over a stream of edges into a near-instant lookup.

With path compression + union by rank, both operations are effectively O(1) amortized (technically the inverse Ackermann α(n)). Memorize the template.

Template

class DSU:
    def __init__(self, n):
        self.p = list(range(n))
        self.r = [0] * n

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

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb: return False
        if self.r[ra] < self.r[rb]: ra, rb = rb, ra
        self.p[rb] = ra
        if self.r[ra] == self.r[rb]: self.r[ra] += 1
        return True

class DSU {
  p: number[]; r: number[];
  constructor(n: number) {
    this.p = [...Array(n).keys()];
    this.r = new Array(n).fill(0);
  }
  find(x: number): number {
    while (this.p[x] !== x) { this.p[x] = this.p[this.p[x]]; x = this.p[x]; }
    return x;
  }
  union(a: number, b: number): boolean {
    let ra=this.find(a), rb=this.find(b);
    if (ra===rb) return false;
    if (this.r[ra]<this.r[rb]) [ra,rb]=[rb,ra];
    this.p[rb] = ra;
    if (this.r[ra]===this.r[rb]) this.r[ra]++;
    return true;
  }
}

Count connected components with union-find:

→ Going deeper: Union-find is the data structure inside Kruskal. See Minimum spanning tree.