Union-Find (Disjoint Set)

Union find maintains a partition of the elements into disjoint sets, and supports two operations: finding which set an element belongs to, and merging two sets. Store the sets as a forest of parent pointers, where the root of each tree is the representative of that set.

Two optimisations make it fast. Path compression flattens the tree as you search, so every node you pass through points directly at the root. Union by rank always attaches the shallower tree under the deeper one, so the trees stay shallow.

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

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

    def union(self, a, b):
        root_a = self.find(a)
        root_b = self.find(b)
        if root_a == root_b:
            return False
        if self.rank[root_a] < self.rank[root_b]:
            root_a, root_b = root_b, root_a
        self.parent[root_b] = root_a
        if self.rank[root_a] == self.rank[root_b]:
            self.rank[root_a] += 1
        return True


uf = UnionFind(5)
print(uf.union(0, 1))
print(uf.union(1, 2))
print(uf.union(0, 2))
print(uf.find(0) == uf.find(2))

The outputs are True, True, False and True. The first two unions merged two different sets, the third found that 0 and 2 were already together and did nothing, and the last check confirms that 0 and 2 are in the same set.

The mechanism is the parent array. A node is a root when its parent is itself, and finding the representative of a set means walking up the parent chain until you reach a root. Path compression rewrites the nodes you pass to point straight at the root, so future finds are shorter, and union by rank keeps the chain lengths logarithmic in the first place.

The trade-off to name is that union find is a one way structure. You can merge sets but you cannot split them, and it only answers questions about membership of a set, not about the graph structure itself. That is exactly what makes it the right tool for counting components or detecting cycles in an undirected graph, and the wrong tool for anything that needs paths or ordering.

Start with every node in its own set. Each edge that joins two different sets reduces the component count by one, so the number of components is the number of nodes minus the number of successful unions.

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

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

    def union(self, a, b):
        root_a = self.find(a)
        root_b = self.find(b)
        if root_a == root_b:
            return False
        if self.size[root_a] < self.size[root_b]:
            root_a, root_b = root_b, root_a
        self.parent[root_b] = root_a
        self.size[root_a] += self.size[root_b]
        return True

    def components(self):
        return len({self.find(node) for node in range(len(self.parent))})


uf = UnionFind(5)
uf.union(0, 1)
uf.union(1, 2)
print(uf.components())

The output is 3. After the two unions the sets are zero, one and two together, node three alone, and node four alone.

The mechanism is that a union of two different sets is exactly one edge that joins two components, so the count goes down by one, while a union inside an existing set changes nothing. The components method counts the distinct roots, which is the number of sets the structure currently holds.

The trade-off to name is the size based union shown here instead of the rank based one. Tracking the size of each set is a slightly simpler way to guarantee the logarithmic bound, because attaching the smaller set under the larger one keeps the depth low. Compared with a graph traversal, union find uses less memory because it does not store the edges at all, but it cannot answer questions about paths or adjacency, so it is the right structure only when you need set membership and merges.