Union Find Disjoint Set

Union Find Disjoint Set (UFDS) is a data structure that keeps track of a set of elements partitioned into a number of disjoint (non-overlapping) subsets. A union-find algorithm is an algorithm that performs two useful operations on such a data structure:

class UnionFindDisjointSet:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n
        self.size = [1] * n

    def find(self, x):
        if self.parent[x] == x:
            return x
        else:
            self.parent[x] = self.find(self.parent[x]) # path compression
            return self.parent[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x != y:
            if self.rank[x] < self.rank[y]:
                x, y = y, x
            self.parent[y] = x
            self.size[x] += self.size[y]
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

    def is_same(self, x, y):
        return self.find(x) == self.find(y)
ufdf = UnionFindDisjointSet(10)
ufdf.union(0, 1)
ufdf.union(2, 3)

print(ufdf.is_same(0, 1))
print(ufdf.is_same(2, 3))
print(ufdf.is_same(0, 2))

# let's join 1 and 2
ufdf.union(1, 2)
print(ufdf.is_same(0, 2))
True
True
False
True

Use Cases

Finding connected components in a graph.

In this case, the elements are the vertices of the graph and the subsets are the connected components. Two elements are in the same subset if and only if they are in the same connected component.

Cycle detection in an undirected graph

When adding an edge, check if the two vertices are already in the same set. If they are, then adding this edge will create a cycle.

Maze generation

Start with a grid of disconnected cells. Pick a random cell and connect it to a random neighbor cell that is not already connected. Repeat until all cells are connected.