Union-Find (Disjoint Set Union): Path Compression and Union by Rank Explained

Introduction to Disjoint Set Union (DSU)

The Union-Find data structure, also known as Disjoint Set Union (DSU), is a data structure that tracks a set of elements partitioned into a number of disjoint (non-overlapping) subsets. It provides near-constant-time operations to add new sets, merge existing sets, and determine whether two elements are in the same set.

DSU is heavily used in graph algorithms, image processing, network connectivity, and type inference systems. Understanding its optimizations—Path Compression and Union by Rank—is fundamental for technical interviews and high-performance system design.

Core Operations

Disjoint Set Union supports two primary operations:

1. Find(i): Determine which subset a particular element i is in. This operation returns the representative (or root) of the set containing i.
2. Union(i, j): Join the subset containing element i and the subset containing element j into a single combined subset.

Naive Implementation

In a naive implementation, we can represent sets using trees where each node points to its parent. An array parent stores these pointers.

class NaiveDisjointSet:
    def __init__(self, size):
        self.parent = list(range(size))

    def find(self, i):
        # Traverse upwards until reaching the root
        if self.parent[i] == i:
            return i
        return self.find(self.parent[i])

    def union(self, i, j):
        root_i = self.find(i)
        root_j = self.find(j)
        if root_i != root_j:
            self.parent[root_i] = root_j

The Problem with the Naive Approach

In the worst-case scenario, the tree can degenerate into a linear linked list (e.g., performing sequential unions). When this happens, the find operation takes $O(N)$ time, rendering the data structure no better than a basic array traversal.

To achieve near $O(1)$ performance, we apply two critical optimizations.

Optimization 1: Path Compression

Path compression is a technique used during the find operation to flatten the structure of the tree. Whenever find traverses up to the root, we can point every node along the path directly to the root.

Visualizing Path Compression

Before find(4):
4 -> 3 -> 2 -> 1 -> 0 (Root)

After find(4) (Path Compression updates pointers directly):
4 -> 0, 3 -> 0, 2 -> 0, 1 -> 0

This dramatically shortens subsequent searches, ensuring that future queries on those nodes execute in $O(1)$ time.

Optimization 2: Union by Rank

Union by Rank ensures that smaller trees are always attached beneath the root of taller trees. We maintain a rank array (or tree height estimate) for each node.

• When executing union(i, j), we attach the root of the lower-rank tree to the root of the higher-rank tree.
• If both roots have the same rank, we arbitrarily choose one as the new root and increment its rank by 1.

Production-Ready Implementation

Below is a complete, optimized implementation of Disjoint Set Union in Python incorporating both Path Compression and Union by Rank.

class DisjointSet:
    def __init__(self, n: int):
        # Initially, each element is its own root
        self.parent = list(range(n))
        # Rank represents the upper bound of the height of the tree
        self.rank = [0] * n
        # Number of connected components
        self.count = n

    def find(self, i: int) -> int:
        # Path compression: point node directly to the root during traversal
        if self.parent[i] == i:
            return i
        self.parent[i] = self.find(self.parent[i])
        return self.parent[i]

    def union(self, i: int, j: int) -> bool:
        root_i = self.find(i)
        root_j = self.find(j)

        if root_i == root_j:
            return False # Already in the same set

        # Union by rank: attach smaller rank tree under higher rank tree
        if self.rank[root_i] < self.rank[root_j]:
            self.parent[root_i] = root_j
        elif self.rank[root_i] > self.rank[root_j]:
            self.parent[root_j] = root_i
        else:
            self.parent[root_j] = root_i
            self.rank[root_i] += 1

        self.count -= 1
        return True

    def get_component_count(self) -> int:
        return self.count

Time Complexity Analysis: The Inverse Ackermann Function

With both Path Compression and Union by Rank enabled, the amortized time complexity for any sequence of $M$ operations (union and find) on $N$ elements is:

$$\mathcal{O}(M \alpha(N))$$

Where $\alpha(N)$ is the inverse Ackermann function.

What is the Ackermann Function?

The Ackermann function $A(m, n)$ grows at an astronomically fast rate. Its inverse, $\alpha(N)$, grows so slowly that for any practical input size $N$ encountered in computer science (even $N = 10^{60}$), $\alpha(N) \le 4$.

Consequently, operations execute in near constant time, $\mathcal{O}(1)$ amortized.

Application: Kruskal's Minimum Spanning Tree Algorithm

One of the most famous applications of DSU is Kruskal's Algorithm for finding the Minimum Spanning Tree (MST) of a connected, undirected graph with weighted edges.

Algorithm Steps

1. Sort all edges in non-decreasing order of their weight.
2. Iterate through the sorted edges. For each edge $(u, v)$, check if $u$ and $v$ belong to the same connected component using DSU find.
3. If they do not form a cycle (i.e., find(u) != find(v)), include the edge in the MST and merge their sets using union(u, v).
4. Repeat until the MST contains $V - 1$ edges.

Kruskal's Implementation in Python

class Edge:
    def __init__(self, u, v, weight):
        self.u = u
        self.v = v
        self.weight = weight

def kruskal_mst(num_vertices, edges):
    # Sort edges by weight
    edges.sort(key=lambda x: x.weight)

    ds = DisjointSet(num_vertices)
    mst_weight = 0
    mst_edges = []

    for edge in edges:
        if ds.union(edge.u, edge.v):
            mst_weight += edge.weight
            mst_edges.append(edge)

            # Optimization: Stop if we have V-1 edges
            if len(mst_edges) == num_vertices - 1:
                break

    return mst_weight, mst_edges

Conclusion

Disjoint Set Union is a stellar example of how algorithmic optimizations yield exponential performance gains. By combining Path Compression (flattening trees during traversal) and Union by Rank (maintaining balanced tree heights), DSU achieves near $\mathcal{O}(1)$ performance, making it indispensable for handling dynamic connectivity queries and optimizing algorithms like Kruskal's MST.