How to Implement Disjoint Set Union Find in Python

Implement a Disjoint Set Union-Find data structure using a Python dictionary for parent tracking, with path compression and connectivity checks.

Medium Python 3.9+ Aug 9, 2026 Dictionaries & sets 14 views 0 copies

disjoint-set union-find graph dictionary path-compression

Python code

34 lines

Python 3.9+

class DisjointSet:
    def __init__(self):
        self.parent = {}

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

    def union(self, x, y):
        # Initialize if not present
        if x not in self.parent:
            self.parent[x] = x
        if y not in self.parent:
            self.parent[y] = y
        root_x = self.find(x)
        root_y = self.find(y)
        if root_x != root_y:
            self.parent[root_x] = root_y

    def connected(self, x, y):
        return x in self.parent and y in self.parent and self.find(x) == self.find(y)


if __name__ == "__main__":
    dsu = DisjointSet()
    dsu.union(1, 2)
    dsu.union(2, 3)
    dsu.union(4, 5)
    print("parent dict:", dsu.parent)
    print("1 and 3 connected:", dsu.connected(1, 3))
    print("1 and 4 connected:", dsu.connected(1, 4))
    print("4 and 5 connected:", dsu.connected(4, 5))

Output

stdout

parent dict: {1: 2, 2: 3, 3: 3, 4: 5, 5: 5}
1 and 3 connected: True
1 and 4 connected: False
4 and 5 connected: True

How it works

The parent dictionary maps each element to its parent node. In find, path compression recursively updates each node to point directly to the root, making future lookups O(alpha(n)) almost constant time. union initializes missing keys, finds both roots, and attaches one root to the other. The connected method checks if both elements exist and share the same root. This implementation uses a dict instead of a list, allowing non-contiguous or non-integer elements like strings.

Common mistakes

Forgetting to initialize missing keys before calling find
Not using path compression, leading to O(n) chains
Checking connectivity without verifying both elements exist

Variations

Use union by rank with arrays for contiguous integer elements
Add a `find_recursive` wrapper that handles missing keys gracefully

Real-world use cases

Detecting connected components in social network graphs for friend suggestions.
Grouping rows in a database migration tool to apply schema changes atomically.
Tracking connectivity between network nodes or IoT devices in a monitoring system.

How to Implement Disjoint Set Union Find in Python

Python code

Output

How it works

Common mistakes

Variations

Real-world use cases

More from Dictionaries & sets

Tutorials

Quizzes

Python code

Output

How it works

Common mistakes

Variations

Real-world use cases

More from Dictionaries & sets

Keep learning

Tutorials

Quizzes