Python Program for Union-Find (Disjoint Set)

Puts friends into groups with union-find, where merging two groups and asking if two people share one each take almost constant time.

Union-find, also called a disjoint set or DSU (disjoint set union), keeps items in separate groups and does two things quickly: union merges two groups into one, and find says which group an item is in. If Ann is friends with Ben, Cal with Dee and Ben with Dee, then Ann, Ben, Cal and Dee all end up in one group, though Ann and Dee have never met. Each group is a tree: every item points to a parent, and the item at the top, which points to itself, is the group's root and stands for the whole group. Two items are in the same group when they have the same root. Kruskal's algorithm for the cheapest network joining every node uses it, and so does any program that asks "are these two connected?" over and over.

Program

class UnionFind:
    """Items in groups that can be merged and looked up."""

    def __init__(self, items):
        self.parent = {item: item for item in items}  # all alone
        self.size = {item: 1 for item in items}

    def find(self, item):
        """The root of item's group."""
        if self.parent[item] != item:
            # Path compression: point item straight at the root
            self.parent[item] = self.find(self.parent[item])
        return self.parent[item]

    def union(self, a, b):
        """Merges the groups of a and b, smaller under larger."""
        root_a, root_b = self.find(a), self.find(b)
        if root_a == root_b:
            return
        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]


friends = UnionFind(["Ann", "Ben", "Cal", "Dee", "Eve", "Fay"])
friends.union("Ann", "Ben")
friends.union("Cal", "Dee")
friends.union("Ben", "Dee")
friends.union("Eve", "Fay")
print("Ann and Dee:", friends.find("Ann") == friends.find("Dee"))
print("Ann and Eve:", friends.find("Ann") == friends.find("Eve"))

Open in the editor Download union_find.py

Output

Ann and Dee: True
Ann and Eve: False

How it works

Visualize Step through the program line by line: each step below shows up as its line runs.

1. At first everyone is alone: each item is its own parent, so its own root, and each group has a size of 1. Both are dictionaries keyed by the item.
2. find() follows the parents up to the root, calling itself one step higher each time (recursion). On the way back it points every item it passed straight at the root. That's path compression: the next find() for any of them takes a single step.
3. union() finds the roots of both items. If the roots are the same, the two are in one group already and there's nothing to do.
4. Otherwise the root of the smaller group goes under the root of the larger, and their sizes are added. That's union by size: hanging the smaller tree under the larger keeps the trees shallow, so find() has few steps to take. Only a root's size is used afterwards.
5. Six people start in six groups, and four friendships merge them into two.
6. Two people are in the same group when find() gives them the same root: Ann and Dee are, Ann and Eve aren't.

Each person's parent after each step, with the ones that change in bold:

union("Ben", "Dee") joins the two roots, Ann and Cal, not Ben and Dee themselves. Both groups have 2 people, so Cal goes under Ann. Dee still points at Cal, two steps from the root, until the find("Dee") on the first print() line points it straight at Ann.

Time and space complexity

With both path compression and union by size, any mix of m finds and unions on n items takes O(m α(n)) time, where α is the inverse Ackermann function. It grows so slowly that it's at most 4 for any number of items a computer could hold, so each operation takes close to constant time, on average over many. A single find() can take a few steps more, but it shortens the path for the ones after it. With union by size alone, no tree is more than log₂ n levels deep, so a find() takes O(log n); with neither, a long chain could make it O(n). The two dicts hold one entry per item: O(n) extra space.

Other ways to write it

With lists and numbered items

When the items are numbered 0 to n − 1, lists can replace the dicts. This find() is a loop, not recursion: as it climbs, it points each item at its grandparent (path halving), which is as fast as path compression without having to come back down. union() also returns whether it merged anything, which Kruskal's algorithm uses to skip an edge that would close a loop.

def find(parent, x):
    while parent[x] != x:
        parent[x] = parent[parent[x]]  # skip to the grandparent
        x = parent[x]
    return x

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

parent = list(range(6))
size = [1] * 6
for a, b in [(0, 1), (2, 3), (1, 3), (4, 5)]:
    union(parent, size, a, b)
print(find(parent, 0) == find(parent, 3))
print(find(parent, 0) == find(parent, 4))

Try it yourself

• Add friends.union("Dee", "Fay") before the print() lines. Now everyone is in one group, and both lines print True.
• Print friends.parent before and after friends.find("Dee") to watch path compression point Dee straight at Ann.
• Count the groups: start a counter at the number of people and take 1 off for each union() that merges two groups. It ends at 2.
• List each group: go through the people and add each one to a dict of lists under their root. The result is {'Ann': ['Ann', 'Ben', 'Cal', 'Dee'], 'Eve': ['Eve', 'Fay']}.

More graphs programs

• Python Program for Breadth-First Search (BFS) Intermediate
• Python Program for Depth-First Search (DFS) Intermediate
• Python Program for Dijkstra's Shortest Path Algorithm Advanced
• Python Program for Topological Sort (Kahn's Algorithm) Advanced