How to Implement UnionFind Algorithm in Python

The
UnionFind
Algorithm in Python 
Implement the
UnionFind
Algorithm in Python 
Implementation of
UnionFind
Algorithm With Rank and Path Compression 
Implementation of
UnionFind
Algorithm With Arrays and Size  Conclusion
This tutorial will discuss how to implement the unionfind
algorithm in Python.
The UnionFind
Algorithm in Python
In programming, selecting the optimal data structure is very important to ensure the efficient performance of our code.
The unionfind
data structure is used to note and keep track of a given set of values partitioned into many disjoint subsets that are not overlapping. This data structure is also known as a disjoint subset
.
It supports two operations on the subsets, Find
and Union
. Let us discuss them below.
The Find
operation finds the subset of a given element. It provides the subset representative, typically an item from this set.
The Union
operation merges two subsets. It combines the subsets only if they belong to the same set, and the two subsets then share the representative.
We use two Find
operations to compare the two elements and check if they belong to the same subset. If they have the same representative, they do, and then we perform the Union
operation to merge the two subsets to which the two elements belonged.
The unionfind
algorithm has different applications like finding the minimum spanning tree, detecting cycles in an undirected graph, etc.
Implement the UnionFind
Algorithm in Python
To implement the unionfind
in Python, we use the concept of trees. The tree’s root can act as a representative, and each node will hold the reference to its parent node.
The unionfind
algorithm will traverse the parent nodes to reach the root and combine two trees by attaching their roots.
Example Code:
class uf_ds:
parent_node = {}
def make_set(self, u):
for i in u:
self.parent_node[i] = i
def op_find(self, k):
if self.parent_node[k] == k:
return k
return self.op_find(self.parent_node[k])
def op_union(self, a, b):
x = self.op_find(a)
y = self.op_find(b)
self.parent_node[x] = y
def display(u, data):
print([data.op_find(i) for i in u])
u = [1, 3, 5, 7]
data = uf_ds()
data.make_set(u)
display(u, data)
data.op_union(1, 5)
display(u, data)
data.op_union(7, 1)
display(u, data)
Output:
[1, 3, 5, 7]
[5, 3, 5, 7]
[5, 3, 5, 5]
In this code, we have implemented a UnionFind
data structure using the class uf_ds
. We initialize a dictionary called parent_node
to keep track of the parent node of each element.
The make_set
method initializes the disjoint sets by assigning each element as its own parent. The op_find
method recursively finds the representative (root) of a given element using path compression
.
The op_union
method merges two sets by making the representative of one set the parent of the other. The display
function is used to print the representative of each element in a given set.
To demonstrate the usage, we create a set with elements [1, 3, 5, 7]
, perform union operations (1, 5)
and (7, 1)
, and display the evolving sets. Notice the result after each operation.
However, the time complexity becomes linear in the worstcase scenario for the above implementation. The trees get skewered and are no better than linked lists.
Thus, two optimization techniques have been introduced to improve the unionfind
algorithm. The first involves the Ranking
of the tree to improve the Union
operation, and the second is Path compression
, which will be discussed in the following.
Implementation of UnionFind
Algorithm With Rank and Path Compression
The Ranking
of the tree is the depth of the tree can increase the time complexity of the naïve approach, so this technique ensures that we attach the root of the smaller tree to the bigger tree. This improves the worstcase time complexity of the unionfind
algorithm in Python to almost O(Log(n))
.
Path compression
is another technique used to improve the efficiency of the UnionFind
algorithm in Python. This approach intends to flatten the given tree and improve the Find
operation.
This intends to make the discovered root node the parent of node n
. This removes the need to traverse through the immediate nodes.
All the elements below this compress when node n
is the root of a given subtree.
These two techniques are efficient in improving the time complexity of the unionfind
algorithm, making it even less than (OLog(n))
. They complement each other.
Example Code:
class ufds:
parent_node = {}
rank = {}
def make_set(self, u):
for i in u:
self.parent_node[i] = i
self.rank[i] = 0
def op_find(self, k):
if self.parent_node[k] != k:
self.parent_node[k] = self.op_find(self.parent_node[k])
return self.parent_node[k]
def op_union(self, a, b):
x = self.op_find(a)
y = self.op_find(b)
if x == y:
return
if self.rank[x] > self.rank[y]:
self.parent_node[y] = x
elif self.rank[x] < self.rank[y]:
self.parent_node[x] = y
else:
self.parent_node[x] = y
self.rank[y] = self.rank[y] + 1
def display(u, data):
print([data.op_find(i) for i in u])
u = [1, 3, 5, 7]
data = ufds()
data.make_set(u)
display(u, data)
data.op_union(1, 5)
display(u, data)
data.op_union(7, 1)
display(u, data)
Output:
[1, 3, 5, 7]
[5, 3, 5, 7]
[5, 3, 5, 5]
In this code, we define a class ufds
representing a UnionFind
data structure. We initialize two dictionaries, parent_node
and rank
, to keep track of the parent nodes and ranks of each element, respectively.
The make_set
method initializes the disjoint sets by assigning each element as its own parent with a rank of 0
. The op_find
method recursively finds the representative (root) of a given element using path compression
.
The op_union
method combines two sets by attaching the smaller tree to the root of the larger tree, considering the ranks to maintain balance. The display
function is used to print the representative of each element in a given set.
Lastly, we demonstrate the usage of this UnionFind
data structure by creating a set with elements [1, 3, 5, 7]
, performing union
operations (1, 5)
and (7, 1)
, and displaying the evolving sets.
Implementation of UnionFind
Algorithm With Arrays and Size
The unionfind
algorithm with arrays and size is a data structure that tracks a set of elements partitioned into disjoint sets. The goal of this data structure is to efficiently determine whether two elements are in the same set and to unite two sets.
The implementation typically includes the use of arrays to represent parent nodes and a separate array to track the size of each set. The size optimization helps in the union operation by always attaching the smaller set to the root of the larger set, which helps maintain balance and improves performance.
Example Code:
class UnionFindWithSize:
def __init__(self, n):
self.parent = list(range(n))
self.size = [1] * n
def find(self, p):
if p != self.parent[p]:
self.parent[p] = self.find(self.parent[p])
return self.parent[p]
def union(self, p, q):
root_p = self.find(p)
root_q = self.find(q)
if root_p != root_q:
if self.size[root_p] < self.size[root_q]:
self.parent[root_p] = root_q
self.size[root_q] += self.size[root_p]
else:
self.parent[root_q] = root_p
self.size[root_p] += self.size[root_q]
n = 10
uf = UnionFindWithSize(n)
uf.union(0, 1)
uf.union(2, 3)
uf.union(4, 5)
uf.union(6, 7)
uf.union(8, 9)
uf.union(1, 9)
print(uf.find(0))
print(uf.find(3))
print(uf.find(5))
print(uf.find(7))
print(uf.find(9))
Output:
0
2
4
6
0
In this code, we’ve implemented the UnionFind
algorithm with size optimization using the UnionFindWithSize
class. We initialize a UnionFind
data structure with 10
elements.
The union
method performs union operations, attaching smaller sets to the root of larger sets based on their sizes. The find
method efficiently determines the representative (root) of a given element, applying path compression to optimize future operations.
We create an instance of the class, perform several union operations, and then execute the find operations on different elements. The output reflects the representatives of elements 0, 2, 4, 6, 8
after union operations, showcasing the efficiency of the UnionFind
algorithm with size optimization in maintaining balanced sets.
Conclusion
This tutorial covered the implementation of the UnionFind
algorithm in Python, emphasizing the importance of optimizing its performance. It introduced techniques like Ranking
and Path Compression
and showcased an implementation with arrays and size for enhanced efficiency.
The provided examples illustrated the algorithm’s versatility in solving problems such as detecting cycles in graphs.
Manav is a IT Professional who has a lot of experience as a core developer in many live projects. He is an avid learner who enjoys learning new things and sharing his findings whenever possible.
LinkedIn