Optimizing Disjoint Set Union

1. The Core Idea: Managing Disjoint Sets

At its heart, a DSU tracks a collection of elements partitioned into a number of non-overlapping (disjoint) sets. It provides three fundamental operations:

1. make_set(v): Creates a new set containing only the element v.
2. find(v): Determines the representative (or "leader") of the set that element v belongs to. If find(a) and find(b) return the same representative, a and b are in the same set.
3. union(a, b): Merges the two sets containing elements a and b into a single set.

To get a strong intuition for what a DSU is and how it works, let's start with a clear visual explanation.

Metaphor: Social Circles

Imagine you're tracking social circles at a large conference.

• make_set: Each person starts as their own social circle of one.
• find: To see if two people, Alice and Bob, are in the same social circle, you find the "most well-known person" (the representative) in each of their groups. If they trace their connections back to the same person, they are in the same circle.
• union: When two people from different circles become friends (e.g., Alice introduces her circle to Charlie's), their two distinct social circles merge into one larger one.

2. A Naive Implementation and Its Pitfall

The simplest way to implement a DSU is by using an array, parent, where parent[i] stores the parent of element i. The representative of a set is an element r such that parent[r] == r.

Technical Explanation

• Initialization: We start with n elements, each in its own set. We can represent this with an array where parent[i] = i for all i.
• find(i): To find the representative of i, we traverse up the chain of parents: i, parent[i], parent[parent[i]], and so on, until we reach a node that is its own parent.
• union(i, j): We find the representatives of i and j, let's call them root_i and root_j. If they are different, we merge the sets by setting one root as the parent of the other, e.g., parent[root_j] = root_i.

The Pitfall: This naive approach can be very slow. If we repeatedly perform unions in an unlucky order, we can create a tree that is just a long, spindly chain.

Caption: A worst-case scenario for a naive DSU. Performing union(0,1), then union(1,2), then union(2,3) creates a chain. A find(3) operation would have to traverse 3 steps to reach the root, 0. In a set of size n, this can take $O(n)$ time.

In this worst-case scenario, the find operation has a time complexity of $O(n)$, which is no better than searching through a linked list.

Python Implementation (Naive)

class DSU_Naive:
    def __init__(self, n):
        # Each element is its own parent initially
        self.parent = list(range(n))

    def find(self, i):
        # Traverse up until we find the root
        while i != self.parent[i]:
            i = self.parent[i]
        return i

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

To make this data structure practical, we need to keep our trees flat. This leads to our first key optimization.

3. Optimization 1: Union by Rank

The problem with the naive union is that it merges sets arbitrarily. Union by Rank (or its close cousin, Union by Size) provides a simple heuristic: when merging two trees, always attach the shorter tree to the root of the taller tree. This minimizes the resulting tree's height.

Technical Explanation

We maintain an extra array, rank, which stores an upper bound on the height of the tree rooted at i.

• Initialization: All ranks start at 0.
• union(i, j):
  1. Find the roots root_i and root_j.
  2. If rank[root_i] < rank[root_j], make root_j the parent of root_i. The ranks don't change.
  3. If rank[root_i] > rank[root_j], make root_i the parent of root_j. The ranks don't change.
  4. If rank[root_i] == rank[root_j], it doesn't matter which way we merge. Pick one (e.g., make root_i the parent of root_j) and increment the rank of the new root (rank[root_i]++).

This strategy ensures that the height of any tree is at most $O(\log n)$. Consequently, the find operation's time complexity improves to $O(\log n)$.

Metaphor: Smart Corporate Mergers

In our company hierarchy metaphor, a naive union is like making a large company's CEO report to a small startup's CEO. This creates a needlessly tall and inefficient structure. Union by Rank is the sensible strategy: when two companies merge, the CEO of the smaller company (in terms of management layers) reports to the CEO of the larger one. This keeps the overall chain of command as short as possible.

4. Optimization 2: Path Compression

Union by Rank gives us a good logarithmic performance guarantee, but we can do even better. Path Compression is a powerful optimization applied to the find operation that dramatically flattens the tree structure over time.

Technical Explanation

The idea is simple: whenever we call find(i), we trace a path from i to its root. After we find the root, we can go back along that path and make every node we visited point directly to the root.

A particularly elegant way to implement this is with recursion:

# In the find method...
def find(self, i):
    if self.parent[i] == i:
        return i
    # Recursively find the root and...
    # ...set the parent of i directly to the root.
    self.parent[i] = self.find(self.parent[i])
    return self.parent[i]

Each time you search for an element's root, you are effectively "caching" the result for all elements along that path, making future searches on them almost instantaneous.

5. The Fully Optimized DSU: A Near-Constant Time Wonder

When we combine Union by Rank with Path Compression, the DSU data structure becomes astonishingly efficient.

Technical Explanation and Complexity

The amortized time complexity of the find and union operations becomes $O(\alpha (n))$, where $\alpha (n)$ is the inverse Ackermann function.

This function grows so slowly that for any practical input size n you could ever encounter (even values larger than the number of atoms in the universe), $\alpha (n)$ is never greater than 5. For all intents and purposes, the optimized DSU offers near-constant time operations. This is a profound result in algorithm analysis.

To understand the theory and implementation details more deeply, the cp-algorithms article is an excellent reference.

Final Python Implementation

Here is the complete, optimized DSU class. This is the implementation you should aim to master for interviews.

class DSU:
    """
    An implementation of Disjoint Set Union (DSU) with Path Compression and Union by Rank.
    """
    def __init__(self, n: int):
        # parent[i] = parent of item i
        self.parent = list(range(n))
        # rank[i] = rank of subtree rooted at i
        self.rank = [0] * n
        # Optional: track number of disjoint sets
        self.count = n

    def find(self, i: int) -> int:
        """
        Finds the representative of the set containing item i with path compression.
        Time Complexity: Amortized O(α(n))
        """
        if self.parent[i] == i:
            return i
        # Path compression: set parent to the root
        self.parent[i] = self.find(self.parent[i])
        return self.parent[i]

    def union(self, i: int, j: int) -> bool:
        """
        Merges the sets containing items i and j using union by rank.
        Returns True if a merge occurred, False otherwise.
        Time Complexity: Amortized O(α(n))
        """
        root_i = self.find(i)
        root_j = self.find(j)

        if root_i != root_j:
            # Union by rank
            if self.rank[root_i] > self.rank[root_j]:
                self.parent[root_j] = root_i
            elif self.rank[root_i] < self.rank[root_j]:
                self.parent[root_i] = root_j
            else:
                self.parent[root_j] = root_i
                self.rank[root_i] += 1
            
            self.count -= 1 # Decrement component count
            return True
        return False

Real-World Application: Counting Connected Components

A classic interview problem and real-world application of DSU is counting the number of connected components in a graph. This is essential in network analysis, image processing (segmenting pixels), and is a key subroutine in algorithms like Kruskal's for finding Minimum Spanning Trees.

The algorithm is straightforward:

1. Initialize a DSU with n nodes, where count = n.
2. Iterate through each edge (u, v) in the graph.
3. For each edge, call dsu.union(u, v).
4. The union function will automatically handle decrementing the count only when two previously disconnected components are merged.
5. After processing all edges, dsu.count gives the number of connected components.

This video from NeetCode demonstrates this exact problem, showing how naturally DSU solves it.

Conclusion

Today we've built the Disjoint Set Union data structure from the ground up, moving from a simple but flawed concept to a highly optimized tool used to solve complex connectivity problems with remarkable efficiency.

Key Takeaways:

• Purpose: DSU is designed to efficiently manage and merge disjoint sets, primarily to answer connectivity questions.
• Naive Implementation: A simple parent array leads to $O(n)$ performance in the worst case, as trees can degenerate into long chains.
• Union by Rank: A crucial union heuristic that keeps trees balanced by attaching shorter trees to taller ones, guaranteeing $O(\log n)$ performance.
• Path Compression: A powerful find optimization that flattens tree structures on the fly by re-parenting nodes directly to the root.
• Combined Power: When used together, these two optimizations yield a near-constant amortized time complexity of $O(\alpha (n))$, making DSU one of the most efficient data structures for its specific task.

Preview of the Next Lesson:
We've now seen how to group items with hash tables and track connectivity with DSU. Next, we will tackle a different kind of problem: answering questions about ranges of data. We will implement the Segment Tree, a powerful hierarchical data structure that can calculate the sum, minimum, maximum, or other properties of an array segment (e.g., from index i to j) in highly efficient logarithmic time.