DSU Optimizations

Two ideas. Independently, each speeds DSU up. Together, they produce one of the most elegant amortized-complexity results in computer science.

Optimization 1 — union by rank

The naive union always re-parents the same way (parent[rx] = ry). That’s how we got the degenerate chain.

Union by rank: track the approximate height of each tree, and always attach the shorter tree under the taller root. Heights then grow only when two equal-height trees are joined — and the new height is just one more.

def union(x, y):
    rx, ry = find(x), find(y)
    if rx == ry:
        return False
    if rank[rx] < rank[ry]:
        rx, ry = ry, rx        # swap so rx is the taller (or equal)
    parent[ry] = rx            # attach shorter under taller
    if rank[rx] == rank[ry]:
        rank[rx] += 1          # equal heights merged → height grows by 1
    return True

rank[] is initialized to all zeros.

The bound: with union by rank alone, every tree’s height is at most log₂(n), so find is O(log n). That’s already a huge improvement over O(n).

Union by size (the alternative)

Same idea, but track the size (element count) of each tree instead of an abstract rank, and attach the smaller-size tree under the larger:

def union(x, y):
    rx, ry = find(x), find(y)
    if rx == ry:
        return False
    if size[rx] < size[ry]:
        rx, ry = ry, rx
    parent[ry] = rx
    size[rx] += size[ry]

Same asymptotic bound (O(log n) height), and size[] is sometimes useful for other reasons (querying component size). In interviews, either is fine; rank is slightly tighter in theory, size is more readable.

Optimization 2 — path compression

The other side of the attack: flatten trees during find.

Every time find(x) walks x → p → pp → ... → root, re-point every node on that walk directly to the root. Next time you find(x) (or any node on that walk), it’s O(1).

def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])   # recursively find root AND re-point
    return parent[x]

The recursive form is the simplest: find(parent[x]) returns the root, and we assign that root as x’s new parent on the way back up. Every node on the walk gets re-parented to the root in a single sweep.

The iterative version is uglier but avoids deep recursion (which can blow the stack for very tall trees in Python/Java):

def find(x):
    # find the root
    root = x
    while parent[root] != root:
        root = parent[root]
    # second pass: re-point everyone on the path
    while parent[x] != root:
        parent[x], x = root, parent[x]
    return root

Both flatten the path. Use whichever you can write cleanly under pressure.

Both, together

struct DSU {
    vector<int> parent, rank_;
    int components;
    DSU(int n) : parent(n), rank_(n, 0), components(n) {
        iota(parent.begin(), parent.end(), 0);
    }
    int find(int x) {
        if (parent[x] != x) parent[x] = find(parent[x]); // path compression
        return parent[x];
    }
    bool unite(int x, int y) {
        int rx = find(x), ry = find(y);
        if (rx == ry) return false;
        if (rank_[rx] < rank_[ry]) swap(rx, ry);
        parent[ry] = rx;
        if (rank_[rx] == rank_[ry]) rank_[rx]++;
        components--;
        return true;
    }
    bool connected(int x, int y) { return find(x) == find(y); }
};

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n
        self.components = n
 
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])  # path compression
        return self.parent[x]
 
    def union(self, x, y):
        rx, ry = self.find(x), self.find(y)
        if rx == ry:
            return False
        if self.rank[rx] < self.rank[ry]:
            rx, ry = ry, rx
        self.parent[ry] = rx
        if self.rank[rx] == self.rank[ry]:
            self.rank[rx] += 1
        self.components -= 1
        return True
 
    def connected(self, x, y):
        return self.find(x) == self.find(y)

class DSU {
    int[] parent, rank;
    int components;
    DSU(int n) {
        parent = new int[n];
        rank = new int[n];
        components = n;
        for (int i = 0; i < n; i++) parent[i] = i;
    }
    int find(int x) {
        if (parent[x] != x) parent[x] = find(parent[x]);  // path compression
        return parent[x];
    }
    boolean union(int x, int y) {
        int rx = find(x), ry = find(y);
        if (rx == ry) return false;
        if (rank[rx] < rank[ry]) { int t = rx; rx = ry; ry = t; }
        parent[ry] = rx;
        if (rank[rx] == rank[ry]) rank[rx]++;
        components--;
        return true;
    }
    boolean connected(int x, int y) { return find(x) == find(y); }
}

Thirty lines (rough). This is the version to memorize.

The amortized complexity — inverse Ackermann

With both optimizations, the amortized cost per operation is O(α(n)), where α is the inverse Ackermann function. Tarjan proved this in 1975.

Why “amortized”? Individual finds can still walk a path — but path compression shortens the path so much for future operations that, averaged across m operations, the total cost is O(m · α(n)).

The Ackermann function grows so fast it makes 2^n look quaint. Its inverse grows so slowly:

For any input that fits in the observable universe, α(n) ≤ 5. In practice, treat DSU as O(1) amortized per operation.

What if you forget one?

• Only path compression, no union by rank. Amortized O(log n) per operation. Still very fast. Often acceptable for interviews.
• Only union by rank, no path compression. Worst-case O(log n) per operation. Same asymptotic ceiling, but no amortization needed for the bound.
• Neither. O(n) worst case per operation. Don’t.

In an interview, write both. It’s only a few extra lines, and you avoid having to explain why you skipped one.

Quick check

Next: the four canonical applications — Kruskal’s, cycle detection, component counting, and weighted DSU.