Solution: Union-Find Implementation
class UnionFind:
"""
Union-Find with path compression and union by rank
Time Complexity:
- find: O(α(n)) ≈ O(1) amortized
- union: O(α(n)) ≈ O(1) amortized
Space: O(n)
"""
def __init__(self, n):
"""Initialize n disjoint sets"""
self.parent = list(range(n)) # Each node is its own parent initially
self.rank = [0] * n # Height of tree
self.count = n # Number of disjoint sets
def find(self, x):
"""
Find root of x's set with path compression
Makes all nodes on path point directly to root
"""
if self.parent[x] != x:
self.parent[x] = self.find(self.parent[x]) # Path compression
return self.parent[x]
def union(self, x, y):
"""
Merge sets containing x and y
Returns True if union performed, False if already in same set
"""
root_x = self.find(x)
root_y = self.find(y)
if root_x == root_y:
return False # Already in same set
# Union by rank: attach shorter tree under taller tree
if self.rank[root_x] < self.rank[root_y]:
self.parent[root_x] = root_y
elif self.rank[root_x] > self.rank[root_y]:
self.parent[root_y] = root_x
else:
self.parent[root_y] = root_x
self.rank[root_x] += 1
self.count -= 1 # Merged two sets
return True
def connected(self, x, y):
"""Check if x and y are in same set"""
return self.find(x) == self.find(y)
def get_count(self):
"""Return number of disjoint sets"""
return self.count
# Example usage
uf = UnionFind(5) # 5 elements: 0, 1, 2, 3, 4
print(uf.count) # 5 sets initially
uf.union(0, 1) # {0,1}, {2}, {3}, {4}
uf.union(2, 3) # {0,1}, {2,3}, {4}
print(uf.count) # 3 sets
print(uf.connected(0, 1)) # True
print(uf.connected(1, 2)) # False
uf.union(1, 3) # {0,1,2,3}, {4}
print(uf.connected(0, 3)) # True
print(uf.count) # 2 sets
Application: Number of Connected Components
def countComponents(n, edges):
"""
LeetCode 323: Number of Connected Components in Undirected Graph
Args:
n: Number of nodes (0 to n-1)
edges: List of edges [[u, v], ...]
Returns:
Number of connected components
Time: O(E * α(n)) ≈ O(E)
Space: O(n)
"""
uf = UnionFind(n)
for u, v in edges:
uf.union(u, v)
return uf.get_count()
# Test
print(countComponents(5, [[0,1],[1,2],[3,4]])) # 2 components
print(countComponents(5, [[0,1],[1,2],[2,3],[3,4]])) # 1 component
Application: Redundant Connection (Cycle Detection)
def findRedundantConnection(edges):
"""
LeetCode 684: Redundant Connection
Given edges that form a tree + 1 extra edge (creates cycle),
return the edge that can be removed to make it a tree again.
A tree with n nodes has exactly n-1 edges.
The extra edge creates a cycle.
Time: O(n)
Space: O(n)
"""
n = len(edges)
uf = UnionFind(n + 1) # Nodes are 1-indexed
for u, v in edges:
if uf.connected(u, v):
# This edge creates a cycle
return [u, v]
uf.union(u, v)
return []
# Test
print(findRedundantConnection([[1,2],[1,3],[2,3]])) # [2,3]
print(findRedundantConnection([[1,2],[2,3],[3,4],[1,4],[1,5]])) # [1,4]
Application: Graph Valid Tree
def validTree(n, edges):
"""
LeetCode 261: Graph Valid Tree
For a graph to be a valid tree:
1. Must be fully connected (n-1 edges, 1 component)
2. Must have no cycles
Union-Find handles both conditions perfectly.
Time: O(n)
Space: O(n)
"""
# Tree with n nodes must have exactly n-1 edges
if len(edges) != n - 1:
return False
uf = UnionFind(n)
for u, v in edges:
# If nodes already connected, adding edge creates cycle
if not uf.union(u, v):
return False
# After all unions, should have exactly 1 component
return uf.get_count() == 1
# Test
print(validTree(5, [[0,1],[0,2],[0,3],[1,4]])) # True
print(validTree(5, [[0,1],[1,2],[2,3],[1,3],[1,4]])) # False (has cycle)
Advanced: Union-Find with Size Tracking
class UnionFindWithSize:
"""
Extended Union-Find that tracks component sizes
"""
def __init__(self, n):
self.parent = list(range(n))
self.size = [1] * n # Size of each component
self.count = 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, x, y):
root_x, root_y = self.find(x), self.find(y)
if root_x == root_y:
return False
# Always attach smaller to larger
if self.size[root_x] < self.size[root_y]:
self.parent[root_x] = root_y
self.size[root_y] += self.size[root_x]
else:
self.parent[root_y] = root_x
self.size[root_x] += self.size[root_y]
self.count -= 1
return True
def get_size(self, x):
"""Get size of component containing x"""
return self.size[self.find(x)]
def get_max_component_size(self):
"""Get size of largest component"""
return max(self.size[self.find(i)] for i in range(len(self.parent)))
Complexity Analysis
Time Complexity (with both optimizations):
- find: O(α(n)) ≈ O(1)
- union: O(α(n)) ≈ O(1)
- For m operations on n elements: O(m · α(n))
Space Complexity:
- Parent array: O(n)
- Rank/size array: O(n)
- Total: O(n)
When to Use Union-Find vs DFS/BFS
Use Union-Find when:
- Dynamic connectivity (edges added over time)
- Only need to check if connected, not find path
- Building MST (Kruskal's algorithm)
- Offline queries (process all before answering)
Use DFS/BFS when:
- Need to find actual path
- Need to visit all nodes
- Static graph (all edges known upfront)
- Directed graphs with specific traversal order
Key Insights
- Union-Find is perfect for dynamic connectivity
- Path compression makes repeated finds extremely fast
- Union by rank keeps trees balanced
- Combined optimizations give near-constant time
- Cannot efficiently split sets (only merge)
Related Problems
