To recap what we attained with our basic implementation, we managed to write a constant-time findPartition method and a worst-case linear-time merge method.
Now, can we do any better than linear, not just for findPartition, but for all the operations on a disjoint set? Well, turns out that yes, we can!
1. From list to trees
The idea is simple: instead of using lists to represent each partition, we will use trees, as shown in figure 5.7. Each partition is then uniquely identified by the root of the tree associated with the partition. The advantage of trees over lists is that if the tree is balanced, any operation on the tree is logarithmic (instead of linear, as for a list).
When we merge two partitions, we will set one tree root as the child of the root of the other tree; see figure 5.8 to see an example.
This is a huge improvement over the naive solution, because we won’t have to update the partition map for any of the other elements in the partitions merged. Instead, each node in the tree will maintain a link to its parent (we don’t need to save links to children, because they are of no use in this case).
The roots of the tree, as mentioned, uniquely identify each partition. So, when we need to find out which partition an element belongs to, we just retrieve the tree node it’s pointing to and walk up to the root of its tree. In method areDisjoint, if we do the same for both elements and then compare the roots found, we can easily see if two elements belong to the same partition (if and only if the two roots are equal).
So, merging two partitions now requires a constant number of changes, plus the number of look-ups needed to find the two roots. Finding the set to which an element belongs (or seeing if two elements belong to the same partition) requires logarithmic time on average (remember when we introduced trees?[1]) but linear time in the worst case. That’s because when merging partitions, we might get unlucky with the choice of which tree’s root we set as child of the other. By randomly choosing each time which root is going to be used as a child of the other, we can make this worst- case scenario unlikely . . . but it would still be possible (although extremely unlucky) to find ourselves in an edge situation such as the one depicted in figure 5.9. This means that our worst-case scenario for merge still requires O(n) lookups and, what’s worst, now even findPartition has linear running time.
Before seeing how we can improve this further, let’s check out some code for our improved version.
2. Implementing the tree version
Let’s see in detail how this improved implementation works. Most code remains unchanged from the previous section, so we won’t show it here. In the methods that have changes with respect to the naive version, we will underline those changes to help readers quickly compare the two versions.
Instead of mapping to an actual set, elements in the partition map point to each element’s parent in the tree. That’s why, as you can see in the book’s repo, we can rename our partitionsMap to parentsMap to make its purpose explicit.
At initialization, we conveniently set each element as its own parent. Trust me on that one; we’ll see why later.
The same change applies to the add method, which otherwise stays unchanged.
The findPartition method (described in listing 5.6) needs quite a bit of tuning to work properly. Two notes on its implementation:
- With respect to the basic implementation, in this case we won’t return a list anymore, but rather the element at the root of the partition’s tree.
- The return value of findPartition might not immediately make sense to an external caller, and in fact this method will mostly be used internally when called by merge and areDisjoint methods.
After getting the element’s parent, we check to see if it’s the element itself. If an element is its own parent, then we know this means we’ve reached the root of the partition’s tree, because of the way we initialize this field, and because, in merge, we never change a root’s parent.
Otherwise, if the current element does have a parent, we walk up the tree toward its root and perform a recursive call to findPartition, returning its result (line #5).
This new implementation of findPartition, as we already mentioned, is not running in constant time anymore. We will have as many recursive calls as the height of the partition’s tree. Since we can’t make any assumption about the trees so far, this means that we possibly have a number of calls proportional to the number of elements in the Universe U, although this is the worst case and, on average, we can expect far better performance.
It might seem so far that we have only made our data structure’s performance worse. We need to define our new implementation of the merge operation, shown in listing 5.7, to see the advantage of using trees.
By comparing the two implementations you can immediately see that this is simpler, although only the last few lines changed. The good news is that we no longer need to iterate through a list of elements! To merge two partitions, we simply need to get to both trees’ roots, and then set one root as the parent of the other. We still need to find those tree roots, though.
The new line we added only requires constant time, so the method runtime is dominated by the two calls to findPartition. As we have seen, they require time proportional to the height of the tree they are called on, and in the worst case this can still be linear. However, in the average case, and especially in the early stages after initialization, we know the height of the trees will be much smaller.
So, in summary, with this implementation we have a disjoint set for which all the operations still require linear time in the worst case, but on average will only need logarithmic time—even for those operations that were constant-time in the naive implementation. Admittedly, that doesn’t sound like a great result if we focus on worst cases. Nevertheless, if you look at it from a different perspective, we’ve already managed to have a more balanced set of operations on our disjoint set, which is especially nice in those contexts where merge is a common operation (while in read-intensive applications, where merge is only executed rarely, the naive implementation could overall be preferable).
Just read through the next section before dismissing the tree solution; it will be worth it.
Source: Rocca Marcello La (2021), Advanced Algorithms and Data Structures, Manning Publications (2021)