Implementing DSU with Path Compression and Union by Rank/Size

1. What is a Disjoint Set Union (DSU)?

At its core, DSU is a data structure that keeps track of a collection of elements partitioned into a number of non-overlapping (disjoint) sets. Think of it like managing social groups. You have a bunch of people, and each person belongs to exactly one group. The DSU helps you quickly figure out which group someone belongs to and lets you merge two groups into one (for example, when two different clubs decide to join forces).

To get a clear visual and conceptual understanding of DSU and its three main operations (make_set, find, and union), let's start with a video.

As the video explains, we can represent our sets as a forest of trees. Each tree corresponds to one set, and the root of the tree is the representative of that set. We implement this using a parent array, where parent[i] stores the parent of element i. If an element is its own parent (parent[i] == i), it is the root of its tree.

The operations are:

• make_set (Initialization): Every element starts in its own set, so we initialize parent[i] = i for all elements i.
• find(i): To find the representative of the set containing i, we travel up the tree by following parent pointers until we reach the root.
• union(i, j): To merge the sets containing i and j, we first find their respective roots. If they are different, we make one root the child of the other.

2. The Problem with a Naive Implementation

The simple approach described above works, but it can be very inefficient. If we are not careful about how we perform the union operation, we can end up creating very tall, skinny trees that resemble linked lists.

When this happens, a find operation has to traverse this long chain, making its time complexity linear, or $O(N)$, where $N$ is the number of elements. For competitive programming, $O(N)$ per query is often too slow.

To solve this, we introduce two powerful optimizations: Union by Rank/Size and Path Compression.

3. Optimization 1: Union by Rank or Size

To prevent our trees from becoming too tall, we need a smarter way to merge them. The idea is simple: when uniting two trees, we should attach the shorter tree to the root of the taller tree. This heuristic helps keep the overall height of the resulting tree from growing unnecessarily.

We can track the "tallness" in two ways:

• Union by Rank: We maintain a rank array. The rank is an upper bound on the height of the tree. When merging, we attach the root with the lower rank to the root with the higher rank. If the ranks are equal, we attach one to the other and increment the rank of the new root.
• Union by Size: We track the number of elements (size) in each set. When merging, we attach the smaller set to the larger one.

Both strategies are effective. We will focus on Union by Rank for our implementation, as it's a classic approach.

4. Optimization 2: Path Compression

Path compression is an optimization applied during the find operation. It's a clever way to flatten the tree structure on the fly.

When we call find(x), we traverse a path from x up to the root. Path compression says: once we've found the root, let's go back down that same path and make every node we visited a direct child of the root.

The next time we call find on any of those nodes, the lookup will be much faster because the path to the root is now just one step.

5. The Optimized Implementation in Python

Now, let's combine these two optimizations to create an efficient DSU implementation. This is the standard template you'll use in competitive programming.

Let's break down that implementation:

class UnionFind:
    def __init__(self, n):
        # Each node is its own parent initially
        self.parent = list(range(n)) 
        # Rank of each tree is initially 1 (or 0, both are valid starting points)
        self.rank = [1] * n 

    def find(self, x):
        # If x is not the root of its set...
        if self.parent[x] != x:
            # ...recursively find the root and set it as x's parent.
            # This is path compression!
            self.parent[x] = self.find(self.parent[x]) 
        return self.parent[x]

    def union(self, x, y):
        # Find the roots of the sets containing x and y
        root_x = self.find(x)
        root_y = self.find(y)
        
        # If they are not already in the same set...
        if root_x != root_y:
            # ...merge them based on rank (Union by Rank).
            if self.rank[root_x] > self.rank[root_y]:
                self.parent[root_y] = root_x
            elif self.rank[root_x] < self.rank[root_y]:
                self.parent[root_x] = root_y
            else:
                # If ranks are equal, one becomes the new root
                # and its rank increases.
                self.parent[root_y] = root_x
                self.rank[root_x] += 1
            return True # Indicates a successful merge
        return False # Indicates they were already in the same set

• __init__: Sets up the parent and rank arrays. It takes $O(N)$ time.
• find(x): The recursive one-liner self.parent[x] = self.find(self.parent[x]) is the magic of path compression. It finds the root and updates the parent of x in a single motion.
• union(x, y): Implements the union by rank logic. It finds the roots, compares their ranks, and attaches the lower-ranked tree to the higher-ranked one. Notice it also handles the case where ranks are equal.

Conclusion

You have now learned how to implement one of the most efficient data structures for handling dynamic connectivity. When combined, path compression and union by rank/size make the DSU operations incredibly fast. 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. For any practical input size in programming competitions, you can consider this to be $O(1)$.

Key Takeaways:

• DSU is used to manage a collection of disjoint sets.
• It is implemented using a parent array to represent a forest of trees, where each tree is a set.
• A naive implementation can lead to $O(N)$ performance for find operations.
• Union by Rank/Size keeps the trees balanced and shallow during union operations.
• Path Compression flattens the trees during find operations.
• Together, these optimizations give DSU nearly constant-time performance for its main operations.

Up Next:
Now that you have built this powerful tool, it's time to use it! In the next lesson, you will apply the DSU data structure to solve connectivity problems, starting with a classic LeetCode problem: finding the number of connected components (or "provinces") in a graph.