-
Notifications
You must be signed in to change notification settings - Fork 0
/
UnionFindA.py
75 lines (63 loc) · 2.37 KB
/
UnionFindA.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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
"""Union-find data structure."""
# based on https://www.ics.uci.edu/~eppstein/PADS/UnionFind.py
class UnionFind:
"""Union-find data structure."""
def __init__(self, n):
"""Create a new empty union-find structure."""
# we have singletons
self.array = list(range(n))
self.size = n
# We give the clusters as a dict
self.clusters = {i: {i} for i in range(n)}
def __getitem__(self, element):
"""Find and return the name of the set containing the element."""
return self.array[element]
def union(self, group1, group2):
"""Find the sets containing the objects and merge them all."""
if len(self.clusters[group1]) > len(self.clusters[group2]):
normal_order = True
to_extend = group1
to_delete = group2
else:
normal_order = False
to_extend = group2
to_delete = group1
# update elements
for i in self.clusters[to_delete]:
self.array[i] = to_extend
# updating the clusters
self.clusters[to_extend].update(self.clusters[to_delete])
del self.clusters[to_delete]
# which cluster is bigger?
return normal_order
def get_cluster(self, group):
"""List of element with id group."""
if group in self.clusters:
if self.clusters[group] == set():
print("Something went wrong!")
return self.clusters[group]
else:
return None
def move(self, elem, group):
"""Move element into group."""
elem_group = self.array[elem]
self.clusters[elem_group].remove(elem)
if self.clusters[elem_group] == set():
del self.clusters[elem_group]
self.array[elem] = group
if group in self.clusters:
self.clusters[group].add(elem)
else:
self.clusters[group] = {elem}
def escape(self, elem, lower, upper):
"""Move element somewhere else."""
elem_group = self.array[elem]
self.clusters[elem_group].remove(elem)
for i in range(lower, upper):
if i not in self.clusters: # if the cluster is empty
self.array[elem] = i
self.clusters[i] = {elem}
break
def __repr__(self):
"""To print only."""
return str(self.array)+str(self.clusters)