-
Notifications
You must be signed in to change notification settings - Fork 77
/
quick_find_union_find.py
56 lines (47 loc) · 1.58 KB
/
quick_find_union_find.py
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
# coding: utf-8
"""
Union-Find (Disjoint Set)
https://en.wikipedia.org/wiki/Disjoint-set_data_structure
"""
class QuickFindUnionFind:
def __init__(self, union_pairs=()):
self.num_groups = 0
self.auto_increment_id = 1
self.element_groups = {
# element: group_id,
}
for p, q in union_pairs:
self.union(p, q)
def __len__(self):
return self.num_groups
# O(1)
def make_group(self, element):
# Initially, every element is in its own group which contains only itself.
group_id = self.element_groups.get(element)
if group_id is None:
# Group id could be arbitrary as long as each group has an unique one.
group_id = self.auto_increment_id
self.element_groups[element] = group_id
self.num_groups += 1
self.auto_increment_id += 1
return group_id
# O(1)
def find(self, p):
try:
return self.element_groups[p]
except KeyError:
# We implicitly create a new group for the new element `p`.
return self.make_group(p)
# O(n)
def union(self, p, q):
p_group_id = self.find(p)
q_group_id = self.find(q)
if p_group_id != q_group_id:
for element, group_id in self.element_groups.items():
# Merge p into q.
if group_id == p_group_id:
self.element_groups[element] = q_group_id
self.num_groups -= 1
# O(1)
def is_connected(self, p, q):
return self.find(p) == self.find(q)