How to Implement Union-Find Algorithm in Python

Manav Narula Feb 02, 2024
  1. The Union-Find Algorithm in Python
  2. Implement the Union-Find Algorithm in Python
  3. Implementation of Union-Find Algorithm With Rank and Path Compression
  4. Implementation of Union-Find Algorithm With Arrays and Size
  5. Conclusion
How to Implement Union-Find Algorithm in Python

This tutorial will discuss how to implement the union-find algorithm in Python.

The Union-Find Algorithm in Python

In programming, selecting the optimal data structure is very important to ensure the efficient performance of our code.

The union-find 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 union-find algorithm has different applications like finding the minimum spanning tree, detecting cycles in an undirected graph, etc.

Implement the Union-Find Algorithm in Python

To implement the union-find 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 union-find 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 Union-Find 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 worst-case 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 union-find 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 Union-Find 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 worst-case time complexity of the union-find algorithm in Python to almost O(Log(n)).

Path compression is another technique used to improve the efficiency of the Union-Find 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 union-find 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 Union-Find 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 Union-Find 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 Union-Find Algorithm With Arrays and Size

The union-find 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 Union-Find algorithm with size optimization using the UnionFindWithSize class. We initialize a Union-Find 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 Union-Find algorithm with size optimization in maintaining balanced sets.

Conclusion

This tutorial covered the implementation of the Union-Find 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.

Author: Manav Narula
Manav Narula avatar Manav Narula avatar

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

Related Article - Python Algorithm