Union Find Disjoint Set Data Structure in DSA Python - Build from Scratch

Union Find Disjoint Set Data Structure

📖 Scenario: Imagine you are managing groups of friends in a social app. You want to know if two people are in the same friend group and also be able to join two groups together.

🎯 Goal: Build a simple Union Find Disjoint Set data structure to track friend groups. You will create the sets, add a helper variable, write the union and find operations, and finally check if two people are in the same group.

📋 What You'll Learn

Create a list called parent where each person is their own parent initially

Create a list called rank to help keep the tree flat

Write a function find(x) to find the root parent of x

Write a function union(x, y) to join the sets of x and y

Use find to check if two people are in the same set

💡 Why This Matters

🌍 Real World

Union Find is used in social networks to find friend groups, in network connectivity to check if computers are connected, and in image processing to find connected regions.

💼 Career

Understanding Union Find helps in solving problems related to grouping, clustering, and network connectivity efficiently, which is common in software engineering and data science roles.

1

Create initial parent list for 5 people

Create a list called parent with 5 elements where each element is its own index: [0, 1, 2, 3, 4]

DSA Python

# Create the parent list where each person is their own parent
# Your code here

Hint

Each person starts as their own parent, so the list is just their own indexes.

2

Create rank list to keep trees flat

Create a list called rank with 5 elements all set to 0 to track the depth of each tree

DSA Python

parent = [0, 1, 2, 3, 4]
# Create the rank list with zeros
# Your code here

Hint

Rank helps us keep the tree shallow by attaching smaller trees under bigger ones.

3

Write find and union functions

Write a function find(x) that returns the root parent of x using path compression, and a function union(x, y) that joins the sets of x and y using rank to decide which root becomes parent

DSA Python

parent = [0, 1, 2, 3, 4]
rank = [0, 0, 0, 0, 0]
# Write find(x) function
# Write union(x, y) function
# Your code here

Hint

Use recursion in find to compress paths. In union, attach smaller rank tree under bigger rank tree.

4

Check if two people are in the same set

Use union(0, 1) and union(1, 2) to join sets, then print find(0) == find(2) to check if person 0 and person 2 are in the same group

DSA Python

parent = [0, 1, 2, 3, 4]
rank = [0, 0, 0, 0, 0]

def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]

def union(x, y):
    rootX = find(x)
    rootY = find(y)
    if rootX != rootY:
        if rank[rootX]  rank[rootY]:
            parent[rootY] = rootX
        else:
            parent[rootY] = rootX
            rank[rootX] += 1

# Join sets and check if 0 and 2 are in the same set
# Your code here

Hint

After joining 0 with 1 and 1 with 2, 0 and 2 should be in the same set, so the print should show True.

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