Union Find Disjoint Set Data Structure in DSA Python

Mental Model

Union Find keeps track of which items belong together in groups without mixing them up. It helps quickly find if two items are in the same group.

Analogy: Imagine a classroom where students form friend groups. Each student knows who their group leader is. When two groups become friends, their leaders join, making one bigger group.

1 -> 1
2 -> 2
3 -> 3
4 -> 4
5 -> 5

Each number points to itself as its own group leader.

Dry Run Walkthrough

Input: Start with 5 elements: 1, 2, 3, 4, 5. Perform union(1, 2), union(3, 4), union(2, 3). Then check if 1 and 4 are connected.

Goal: Group elements so that 1, 2, 3, 4 are in one group and 5 alone. Confirm 1 and 4 are connected.

Step 1: union(1, 2): connect groups of 1 and 2

1 -> 1
2 -> 1
3 -> 3
4 -> 4
5 -> 5

Why: We make 1 the leader of 2's group to join them.

Step 2: union(3, 4): connect groups of 3 and 4

1 -> 1
2 -> 1
3 -> 3
4 -> 3
5 -> 5

Why: We make 3 the leader of 4's group to join them.

Step 3: union(2, 3): connect groups of 2 and 3

1 -> 1
2 -> 1
3 -> 1
4 -> 3
5 -> 5

Why: We find leaders of 2 (1) and 3 (3), then make 1 leader of 3's group to join all.

Step 4: find(4): follow leader chain 4 -> 3 -> 1

1 -> 1
2 -> 1
3 -> 1
4 -> 1
5 -> 5

Why: Path compression updates 4's leader directly to 1 for faster future finds.

Step 5: connected(1, 4): check if leaders are same

1 -> 1
2 -> 1
3 -> 1
4 -> 1
5 -> 5

Why: Both have leader 1, so they are connected.

Result:

1 -> 1
2 -> 1
3 -> 1
4 -> 1
5 -> 5

1 and 4 are connected: True

Annotated Code

DSA Python

class UnionFind:
    def __init__(self, size: int):
        self.parent = [i for i in range(size + 1)]  # Each node is its own leader
        self.rank = [0] * (size + 1)  # Rank helps keep tree shallow

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

    def union(self, x: int, y: int) -> None:
        rootX = self.find(x)  # Find leader of x
        rootY = self.find(y)  # Find leader of y
        if rootX != rootY:
            if self.rank[rootX] < self.rank[rootY]:
                self.parent[rootX] = rootY  # Attach smaller tree under bigger
            elif self.rank[rootX] > self.rank[rootY]:
                self.parent[rootY] = rootX
            else:
                self.parent[rootY] = rootX
                self.rank[rootX] += 1  # Increase rank if same

    def connected(self, x: int, y: int) -> bool:
        return self.find(x) == self.find(y)

# Driver code using dry_run input
uf = UnionFind(5)
uf.union(1, 2)
uf.union(3, 4)
uf.union(2, 3)

# Print parent array to show groups
for i in range(1, 6):
    uf.find(i)  # Path compression updates
print(' -> '.join(str(uf.parent[i]) for i in range(1, 6)) + ' -> null')

# Check if 1 and 4 are connected
print('1 and 4 are connected:', uf.connected(1, 4))

if self.parent[x] != x: self.parent[x] = self.find(self.parent[x])

compress path by pointing x directly to leader

rootX = self.find(x) rootY = self.find(y)

find leaders of both elements

if rootX != rootY:

only union if leaders differ to avoid cycles

if self.rank[rootX] < self.rank[rootY]: self.parent[rootX] = rootY

attach smaller rank tree under bigger rank tree

elif self.rank[rootX] > self.rank[rootY]: self.parent[rootY] = rootX

attach smaller rank tree under bigger rank tree

else: self.parent[rootY] = rootX self.rank[rootX] += 1

if ranks equal, attach and increase rank

return self.find(x) == self.find(y)

check if two elements share the same leader

Output

1 -> 1 -> 1 -> 1 -> 5 -> null

1 and 4 are connected: True

Complexity Analysis

Time: O(α(n)) per operation, where α is inverse Ackermann function, almost constant because path compression and union by rank keep trees flat

Space: O(n) for storing parent and rank arrays for n elements

vs Alternative: Naive union without path compression or rank can degrade to O(n) per find, making operations slower on large sets

Edge Cases

Union or find on single element

Element is its own leader; union does nothing if same element

DSA Python

if rootX != rootY:

Union of already connected elements

No change to structure; avoids cycles

DSA Python

if rootX != rootY:

Find on element not connected to others

Returns element itself as leader

DSA Python

if self.parent[x] != x:

When to Use This Pattern

When you need to quickly group items and check if two items belong to the same group, use Union Find because it efficiently merges groups and finds leaders.

Common Mistakes

Mistake: Not using path compression in find causes slow repeated searches
Fix: Add path compression by updating parent[x] during find

Mistake: Not using union by rank causes tall trees and slow operations
Fix: Use rank to attach smaller tree under bigger tree

Mistake: Union without checking if leaders differ causes cycles
Fix: Only union if leaders are different

Summary

It groups elements into sets and quickly checks if two elements are in the same set.

Use it when you need to manage dynamic connectivity or grouping problems.

The key is to keep track of group leaders and flatten the structure for fast queries.

Practice

1. What is the main purpose of the Union Find Disjoint Set data structure?

easy

A. To efficiently track and merge groups of connected elements

B. To sort elements in ascending order

C. To store key-value pairs for fast lookup

D. To implement a queue with fast enqueue and dequeue

Solution

1. Step 1: Understand the purpose of Union Find

   The Union Find data structure is designed to keep track of which elements belong to which groups or sets, especially when these groups merge.

2. Step 2: Compare with other data structures

   Sorting, key-value storage, and queue operations are unrelated to Union Find's main goal.

3. Final Answer:

   To efficiently track and merge groups of connected elements -> Option A

4. Quick Check:

   Union Find = track and merge groups [OK]

Hint: Union Find = manage connected groups fast [OK]

Common Mistakes:

• Confusing Union Find with sorting algorithms
• Thinking it stores key-value pairs
• Mixing it up with queue or stack data structures

2. Which of the following is the correct way to initialize a Union Find parent array for 5 elements in Python?

easy

A. parent = [None] * 5

B. parent = [1, 1, 1, 1, 1]

C. parent = [5, 5, 5, 5, 5]

D. parent = [0, 1, 2, 3, 4]

Solution

1. Step 1: Understand initialization of Union Find

   Each element starts as its own parent, so the parent array should have each index equal to itself.

2. Step 2: Check each option

   parent = [0, 1, 2, 3, 4] correctly sets parent[i] = i for i in 0 to 4. Others do not represent this.

3. Final Answer:

   parent = [0, 1, 2, 3, 4] -> Option D

4. Quick Check:

   Parent array init = self indices [OK]

Hint: Parent array starts with self indices [OK]

Common Mistakes:

• Setting all parents to 1 or a fixed number
• Using None which breaks find operations
• Confusing size array with parent array

3. What is the output of the following code?

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    def union(self, a, b):
        rootA = self.find(a)
        rootB = self.find(b)
        if rootA != rootB:
            self.parent[rootB] = rootA

uf = UnionFind(4)
uf.union(0, 1)
uf.union(2, 3)
uf.union(1, 3)
print(uf.parent)

medium

A. [0, 0, 0, 2]

B. [0, 0, 0, 0]

C. [0, 0, 2, 2]

D. [0, 1, 2, 3]

Solution

1. Step 1: Trace unions step-by-step

   Init: [0,1,2,3]
   Union(0,1): find(0)=0, find(1)=1 -> parent[1]=0 -> [0,0,2,3]
   Union(2,3): find(2)=2, find(3)=3 -> parent[3]=2 -> [0,0,2,2]
   Union(1,3): find(1)=0, find(3)=2 (sets parent[3]=2) -> parent[2]=0 -> [0,0,0,2]

2. Step 2: Verify print output

   No additional find calls before print. parent[3] points to 2 (now root 0), but not fully compressed to 0. Output: [0, 0, 0, 2]

3. Final Answer:

   [0, 0, 0, 2] -> Option A

4. Quick Check:

   parent = [0,0,0,2] after partial compression [OK]

Hint: Trace find calls and compression timing during union [OK]

Common Mistakes:

• Forgetting last union sets parent[2]=0 -> [0,0,2,2]
• Over-applying path compression -> [0,0,0,0]
• Ignoring unions entirely

4. Identify the bug in this Union Find union method:

def union(self, a, b):
    rootA = self.find(a)
    rootB = self.find(b)
    if rootA != rootB:
        self.parent[rootA] = rootB

medium

A. The find method is not called correctly

B. It should assign parent of rootB to rootA, not the other way

C. The if condition should be rootA == rootB

D. The method is missing a return statement

Solution

1. Step 1: Understand union direction

   Union should attach one root to the other to merge sets. Usually, parent of rootB is set to rootA or vice versa consistently.

2. Step 2: Check given code

   Code sets parent[rootA] = rootB, which reverses the usual convention and can cause incorrect tree structure or confusion.

3. Final Answer:

   It should assign parent of rootB to rootA, not the other way -> Option B

4. Quick Check:

   Union sets parent[rootB] = rootA [OK]

Hint: Union sets parent of second root to first root [OK]

Common Mistakes:

• Swapping roots incorrectly in union
• Not calling find before union
• Confusing equality check in if condition

5. Given a graph with 6 nodes and edges [(0,1), (1,2), (3,4)], how many connected components does the Union Find structure report after union operations on these edges?

hard

A. 4

B. 2

C. 3

D. 5

Solution

1. Step 1: Initialize 6 separate nodes

   Initially, each node is its own component: 6 components total.

2. Step 2: Perform unions for edges

   Union(0,1) merges nodes 0 and 1 -> components reduce to 5.
   Union(1,2) merges node 2 with 0-1 group -> components reduce to 4.
   Union(3,4) merges nodes 3 and 4 -> components reduce to 3.

3. Step 3: Count remaining components

   Node 5 remains alone -> total components = 3.

4. Final Answer:

   3 -> Option C

5. Quick Check:

   6 nodes - 3 unions = 3 components [OK]

Hint: Components = nodes - unions done [OK]

Common Mistakes:

• Counting edges instead of components
• Forgetting isolated nodes count as components
• Miscounting after unions