Python/UnionFind.py at main · achapi/Python · GitHub

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
fromcollectionsimportdefaultdict

classUnionFind():
def__init__(self, n):
self.n=n
self.parents= [-1] *n

deffind(self, x):
ifself.parents[x] <0:
returnx
else:
self.parents[x] =self.find(self.parents[x])
returnself.parents[x]

defunion(self, x, y):
x=self.find(x)
y=self.find(y)

ifx==y:
return

ifself.parents[x] >self.parents[y]:
x, y=y, x

self.parents[x] +=self.parents[y]
self.parents[y] =x

defsize(self, x):
return-self.parents[self.find(x)]

defsame(self, x, y):
returnself.find(x) ==self.find(y)

defmembers(self, x):
root=self.find(x)
return [iforiinrange(self.n) ifself.find(i) ==root]

defroots(self):
return [ifori, xinenumerate(self.parents) ifx<0]

defgroup_count(self):
returnlen(self.roots())

defall_group_members(self):
group_members=defaultdict(list)
formemberinrange(self.n):
group_members[self.find(member)].append(member)
returngroup_members

def__str__(self):
return'\n'.join(f'{r}: {m}'forr, minself.all_group_members().items())