Implementing Optimized Disjoint-Set Union

Hello! In our previous lesson, we explored a greedy strategy for solving reachability problems by tracking the "farthest reach." Today, we shift our focus from linear paths to a more complex and common scenario in algorithms: managing connections and groups within a set of elements.

This lesson introduces the Disjoint-Set Union (DSU) data structure, also known as Union-Find. This is a specialized and remarkably efficient tool for tracking a collection of disjoint (non-overlapping) sets. You'll learn to implement it from the ground up, incorporating two powerful optimizations—path compression and union by size—that make its operations almost instantaneous. Mastering DSU is crucial for a wide range of problems involving connectivity, network analysis, and partitioning.

What is a Disjoint-Set Union?

Imagine you have a set of items, and you want to group them together. Initially, each item is in its own group of one. You can then perform two primary operations:

1. Union: Merge two groups into a single, larger group.
2. Find: Determine which group a specific item belongs to.

The DSU data structure is designed to perform these operations with extreme efficiency. It's often used to answer questions like: "Are these two people in the same social network?" or "Does adding this road connect two previously separate parts of a city?"

The following video provides a friendly, high-level introduction using a great real-world analogy of a company's organizational chart.

Representing Disjoint Sets

As the video mentions, the most common way to represent these sets is not with complex objects, but with a simple array, let's call it parent. Each index i in the array corresponds to an element. The value parent[i] stores the parent of element i.

This creates a "forest" of trees, where each tree represents one set. The root of each tree is the representative of that set. An element is a root if it is its own parent (i.e., parent[i] === i).

The find operation works by traversing up the parent pointers from a starting node until it reaches the root. The union operation merges two trees by making the root of one tree a child of the root of the other.

The Problem with a Naive Approach

How we perform the union operation matters a great deal. A naive implementation might arbitrarily pick one root to become the child of the other. This can lead to a serious problem: the trees can become very tall and unbalanced, resembling a linked list.

This inefficiency is exactly the kind of bottleneck we want to avoid. To fix this, we introduce two powerful optimizations.

Optimization 1: Union by Size

To prevent our trees from becoming tall and skinny, we can use a simple heuristic: when merging two trees, always attach the smaller tree to the root of the larger tree. This helps keep the overall tree height low.

We can track the "largeness" of a tree in two ways: by its height (called union by rank) or by the number of nodes it contains (called union by size). Union by size is often more intuitive and is what we'll focus on implementing.

Here's how Union by Size works:

1. Maintain a size array, where size[i] stores the number of elements in the set if i is the root. Initially, all sizes are 1.
2. When uniting two sets with roots rootA and rootB, compare their sizes.
3. If size[rootA] is less than size[rootB], attach rootA to rootB. That is, set parent[rootA] = rootB and update the size of the new combined set: size[rootB] += size[rootA].
4. Otherwise, do the opposite: attach rootB to rootA.

The following image illustrates the difference between a naive union and a union guided by a similar principle, "union by rank." The key takeaway is that the optimized version on the right produces a much flatter, more balanced tree.

Optimization 2: Path Compression

The second optimization dramatically speeds up the find operation. The idea is that after we traverse a path from a node up to its root, we can make all the nodes we visited point directly to that root.

This flattens the tree structure every time we perform a find. Subsequent find operations on any of those nodes (or their children) will then be much faster, often taking just one step.

The most elegant way to implement path compression is using recursion in the find method. The following video demonstrates this idea perfectly.

A Visual Playground

Before we dive into the full implementation, it's incredibly helpful to see these concepts in action. The VisuAlgo website has a fantastic interactive visualization for DSU. I highly recommend spending a few minutes playing with it.

Implementing DSU in TypeScript

Now, let's put everything together into a reusable TypeScript class. This implementation will include the constructor, the find operation with path compression, and the union operation with the union by size optimization, directly fulfilling our lesson's goal.

class DisjointSet {
    private parent: number[];
    private size: number[];
    private numComponents: number;

    constructor(n: number) {
        // Initially, each element is in its own set.
        // The parent of an element is itself.
        this.parent = Array.from({ length: n }, (_, i) => i);
        
        // The size of each set is initially 1.
        this.size = new Array(n).fill(1);

        // The number of disjoint sets is n.
        this.numComponents = n;
    }

    // Find the representative (root) of the set containing 'i'.
    // Applies path compression along the way.
    find(i: number): number {
        if (this.parent[i] === i) {
            return i;
        }
        // Recursively find the root and update the parent pointer.
        this.parent[i] = this.find(this.parent[i]);
        return this.parent[i];
    }

    // Merge the sets containing 'i' and 'j'.
    // Returns true if a merge happened, false otherwise.
    union(i: number, j: number): boolean {
        const rootI = this.find(i);
        const rootJ = this.find(j);

        // If they are already in the same set, do nothing.
        if (rootI === rootJ) {
            return false;
        }

        // Union by size: attach the smaller tree to the root of the larger tree.
        if (this.size[rootI] < this.size[rootJ]) {
            this.parent[rootI] = rootJ;
            this.size[rootJ] += this.size[rootI];
        } else {
            this.parent[rootJ] = rootI;
            this.size[rootI] += this.size[rootJ];
        }
        
        // A merge occurred, so decrement the number of components.
        this.numComponents--;
        return true;
    }

    // Check if elements 'i' and 'j' are in the same set.
    isConnected(i: number, j: number): boolean {
        return this.find(i) === this.find(j);
    }
    
    // Get the number of disjoint sets.
    getComponentCount(): number {
        return this.numComponents;
    }
}

// --- Example Usage ---
const ds = new DisjointSet(10); // Create a DSU with 10 elements (0-9)

ds.union(0, 1);
ds.union(1, 2);
console.log(`0 and 2 are connected: ${ds.isConnected(0, 2)}`); // true

ds.union(5, 6);
console.log(`0 and 5 are connected: ${ds.isConnected(0, 5)}`); // false

console.log(`Initial number of components: 10`);
console.log(`Components after 3 unions: ${ds.getComponentCount()}`); // 7

ds.union(0, 5);
console.log(`0 and 5 are now connected: ${ds.isConnected(0, 5)}`); // true
console.log(`Components after 4 unions: ${ds.getComponentCount()}`); // 6

Time and Space Complexity

The beauty of DSU with both union by size/rank and path compression lies in its performance.

• Time Complexity: The amortized time complexity for both find and union is nearly constant, denoted as $O(\alpha (N))$, where $\alpha (N)$ is the Inverse Ackermann function. This function grows so slowly that for any practical input size you'd encounter in an interview (even for $N$ up to $10^{80}$), its value is less than 5. For all intents and purposes, you can consider these operations to be constant time.
• Space Complexity: The space complexity is $O(N)$ to store the parent and size arrays.

This incredible efficiency makes DSU the go-to data structure for a wide variety of connectivity problems.

Conclusion

In this lesson, we delved into the Disjoint-Set Union data structure, a powerful tool for managing dynamic sets of elements.

Here are the key takeaways:

• Core Idea: DSU tracks a collection of disjoint sets using a parent array to represent a forest of trees.
• Key Operations: find identifies a set's representative (the root), and union merges two sets.
• Essential Optimizations:
  • Union by Size: Keeps trees balanced and shallow by attaching the smaller tree to the root of the larger one.
  • Path Compression: Flattens the tree structure during find operations by making nodes point directly to the root.
• Performance: The combination of these optimizations yields nearly constant-time performance for find and union, making it one of the most efficient data structures available.

You've now successfully implemented a foundational and highly practical data structure. In our next lesson, we will immediately put it to use to solve a classic graph problem: determining whether adding an edge to a graph creates a cycle. This will solidify your understanding of DSU and demonstrate its power in a real-world context.