-
Notifications
You must be signed in to change notification settings - Fork 2
Expand file tree
/
Copy path__init__.py
More file actions
74 lines (62 loc) · 2.53 KB
/
__init__.py
File metadata and controls
74 lines (62 loc) · 2.53 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
class DisjointSetUnion:
"""A class for the Union-Find (Disjoint Set Union) data structure."""
def __init__(self, size: int):
"""Initializes the data structure with 'size' elements, each in its own set."""
if size <= 0:
raise ValueError("Size must be a positive integer.")
self.root = list(range(size))
self.rank = [1] * size # For union by rank
self.count = size # Number of disjoint sets
def find(self, i: int) -> int:
"""Finds the representative (root) of the set containing element 'i'."""
if self.root[i] == i:
return i
# Path compression: make all nodes on the path point to the root
self.root[i] = self.find(self.root[i])
return self.root[i]
def union(self, i: int, j: int) -> bool:
"""
Merges the sets containing elements 'i' and 'j'.
Returns True if a merge occurred, False if they were already in the same set.
"""
root_i = self.find(i)
root_j = self.find(j)
if root_i != root_j:
# Union by rank: attach the smaller tree to the larger tree
if self.rank[root_i] > self.rank[root_j]:
self.root[root_j] = root_i
elif self.rank[root_i] < self.rank[root_j]:
self.root[root_i] = root_j
else:
self.root[root_j] = root_i
self.rank[root_i] += 1
self.count -= 1
return True
return False
def get_count(self) -> int:
"""Returns the current number of disjoint sets."""
return self.count
class UnionFind:
"""A minimal Union-Find data structure with path compression."""
def __init__(self, size: int):
"""Initializes the data structure with 'size' elements."""
if size <= 0:
raise ValueError("Size must be a positive integer.")
self.parent = list(range(size))
def find(self, x: int) -> int:
"""Finds the representative (root) of the set containing element 'x'."""
if self.parent[x] != x:
# Path compression
self.parent[x] = self.find(self.parent[x])
return self.parent[x]
def union(self, x: int, y: int) -> bool:
"""
Merges the sets containing elements 'x' and 'y'.
Returns True if a merge occurred, False if already in same set.
"""
root_x = self.find(x)
root_y = self.find(y)
if root_x != root_y:
self.parent[root_y] = root_x
return True
return False