Union-find: two lines of optimisation make every operation near constant

Union-find answers one question: elements keep getting merged into groups, and at any moment you need to know whether two of them are in the same group. It fits in 30 lines, but it only reaches O(alpha(n)) with the right two lines of optimisation.

The implementation

class Dsu {
  parent: number[];
  size: number[];

  constructor(n: number) {
    this.parent = Array.from({ length: n }, (_, i) => i);
    this.size = new Array(n).fill(1);
  }

  find(x: number): number {
    let root = x;
    while (this.parent[root] !== root) root = this.parent[root];
    while (this.parent[x] !== root) {
      const next = this.parent[x];
      this.parent[x] = root;
      x = next;
    }
    return root;
  }

  union(a: number, b: number): boolean {
    let x = this.find(a);
    let y = this.find(b);
    if (x === y) return false;
    if (this.size[x] < this.size[y]) [x, y] = [y, x];
    this.parent[y] = x;
    this.size[x] += this.size[y];
    return true;
  }
}

The second while inside find is path compression: it hangs the whole path from x to the root directly off the root, so the next lookup takes one step. The if (this.size[x] < this.size[y]) inside union is union by size: attach the smaller tree under the larger one, or single-sided chains keep growing.

Why you want both

Optimisation	Cost on its own
Neither	Worst case O(n) per operation, degenerating into a linked list
Path compression only	Amortised O(log n)
Union by size only	O(log n), and it never compresses
Both	O(alpha(n)), inverse Ackermann, at most 4 in practice

The inverse Ackermann function is below 5 even at astronomical input sizes, which is what almost constant means: for any input you will ever see, it is constant.

Where it gets used

Beyond Kruskal in a textbook, it shows up in counting connected components of an undirected graph, merging adjacent pixels of the same class in image processing, equivalence classes of types in a compiler, and answering whether two nodes are connected while edges arrive online.

Three traps I have hit

• A recursive implementation blows the stack. Recursive path compression reads beautifully and overflows on a chain of a million nodes. Use the two-pass iterative version above.
• Compression makes rank wrong. With union by rank, compression leaves the recorded rank too high because the tree is shorter now. Size stays truthful, so prefer it.
• Only one path gets compressed. The two find calls in union compress the paths from a and b to the root; intermediate nodes are compressed the next time they are queried. That is correct, and no manual walk over the subtree is needed.

The value of union-find is not in its length but in being willing to write those two while loops.