Implementing Union-Find

1. What is Union-Find?

Imagine you have a collection of items, and you want to group them into a number of disjoint (non-overlapping) sets. The Union-Find data structure is designed to do two things very quickly:

1. Find: Determine which set a particular item belongs to.
2. Union: Merge two sets into a single new set.

Its primary application is to solve dynamic connectivity problems. For example, given a set of nodes in a graph, can we efficiently answer queries like "Are node A and node B connected?" while edges are being added to the graph dynamically? A naive DFS/BFS would be too slow if you have many queries. Union-Find provides a near-constant time solution.

To build a strong mental model, let's start with a short video that offers a great analogy and a high-level visual introduction.

As the video explains, we can represent these sets as a forest of trees, where each tree corresponds to a set. The root of the tree acts as the representative (or identifier) for that set.

• The find(x) operation involves traversing from node x up to the root of its tree.
• The union(x, y) operation merges the sets containing x and y by making the root of one tree a child of the root of the other.

2. The Need for Optimizations

A simple implementation of this idea has a critical flaw. If we aren't careful about how we perform the union operation, we can end up creating very tall, skewed trees that resemble linked lists.

Consider uniting elements in a chain: union(0,1), then union(1,2), then union(2,3), and so on. The resulting tree would be a straight line. In this worst-case scenario, a find operation would have to traverse all the way up the chain, making its time complexity $O(N)$. For interview settings where performance is key, this is unacceptable.

To achieve the near-constant time complexity that makes Union-Find so powerful, we need two optimizations.

3. The Key Optimizations: Path Compression and Union by Rank/Size

These two techniques work together to keep the trees in our forest extremely flat, ensuring that the find operation remains fast.

Union by Rank / Union by Size

This optimization applies to the union operation. Instead of arbitrarily attaching one tree to another, we use a heuristic to decide which tree goes under which.

• Union by Rank: We keep track of the "rank" of each root, which is an upper bound on the height of its tree. When uniting two trees, we always attach the root of the shorter tree (lower rank) to the root of the taller tree (higher rank). This prevents the overall tree height from growing unnecessarily. If the ranks are equal, we can pick one, attach the other, and increment the rank of the new root.
• Union by Size: This is a similar, often more intuitive, alternative. We track the number of nodes (size) in the set represented by each root. When uniting, we attach the root of the smaller set to the root of the larger set.

Path Compression

This optimization applies to the find operation. When we call find(x), we traverse the path from x up to its root. Path compression does a simple, brilliant thing: after finding the root, it makes every node along that path a direct child of the root.

With these two optimizations combined, the amortized time complexity of both union and find operations becomes nearly constant, specifically $O(\alpha (N))$, where $\alpha (N)$ is the Inverse Ackermann function. This function grows so slowly that for any practical input size N, it's less than 5.

Let's watch a detailed walkthrough of implementing these optimizations.

4. Implementation Template

Now, let's consolidate everything into a clean, reusable Python class. This template is what you should aim to reproduce in an interview. You will see it implements the logic discussed in the video.

We will have a class UnionFind that is initialized with the number of elements n. It will maintain a parent array and either a rank or size array.

• __init__(self, n): Initializes parent[i] = i for all elements, and sets their initial rank to 0 or size to 1.
• find(self, i): Finds the root of i and performs path compression using recursion.
• union(self, i, j): Merges the sets containing i and j using union by rank or union by size.

The following resource provides excellent, clean code for both variations.

Here's a summary of the Union by Size implementation, which is a very common interview choice:

class UnionFind:
    def __init__(self, n):
        # Each node is its own parent initially
        self.parent = list(range(n))
        # Each component has a size of 1 initially
        self.size = [1] * n

    def find(self, i):
        # Base case: if i is its own parent, it's the root
        if self.parent[i] == i:
            return i
        
        # Recursive call to find the root, with path compression
        self.parent[i] = self.find(self.parent[i])
        return self.parent[i]

    def union(self, i, j):
        # Find the roots of the sets containing i and j
        root_i = self.find(i)
        root_j = self.find(j)
        
        # If they are already in the same set, do nothing
        if root_i != root_j:
            # Union by size: attach the smaller tree to the root of the larger tree
            if self.size[root_i] < self.size[root_j]:
                root_i, root_j = root_j, root_i # Ensure root_i is the larger set
            
            self.parent[root_j] = root_i
            self.size[root_i] += self.size[root_j]
            return True # Return True to indicate a successful merge
        
        return False # Return False if they were already connected

Conclusion

You have now built one of the most efficient data structures for handling dynamic connectivity. Its elegance and performance make it a favorite in competitive programming and a valuable tool for your interview arsenal.

Key Takeaways:

• Union-Find (DSU) is specialized for managing disjoint sets and answering connectivity queries.
• A naive implementation leads to $O(N)$ performance, which is too slow.
• Union by Rank/Size keeps the underlying trees flat during union operations.
• Path Compression flattens the tree structure during find operations.
• Together, these optimizations give Union-Find a nearly constant amortized time complexity ($O(\alpha (N))$).
• The space complexity is $O(N)$ to store the parent and rank/size arrays.

This table from the "Union-Find" chapter in shawonnotes is a great summary of when to use this data structure.

Preview of the Next Lesson:

Now that you have mastered the implementation of the Union-Find data structure, our next lesson will focus on its practical application. We will apply Union-Find to efficiently solve common graph connectivity problems, such as counting the number of connected components in a graph and detecting redundant connections. You'll see just how powerful this tool is in action.