-
Notifications
You must be signed in to change notification settings - Fork 4
/
union_find.py
47 lines (35 loc) · 1.15 KB
/
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
import sys
sys.path.append("..")
"""
Implementation of Union-Find with path compression and Union by Rank
"""
class UnionFind:
def __init__(self, num_nodes):
self.__ranks__= [0 for __ in range(num_nodes + 1)]
self.__leaders__ = [x for x in range(num_nodes + 1)]
self.__size__ = num_nodes
def find(self, u):
leader = self.__leaders__[u]
if leader == u:
return leader
else:
next_leader = self.find(leader)
self.__leaders__[u] = next_leader # path compression
return next_leader
def union(self, u, v):
u_leader, v_leader = self.find(u), self.find(v)
if u_leader == v_leader:
return False
u_rank, v_rank = self.__ranks__[u_leader], self.__ranks__[v_leader]
if v_rank > u_rank:
self.__leaders__[u_leader] = v_leader
else:
self.__leaders__[v_leader] = u_leader
if u_rank == v_rank:
self.__ranks__[u_leader]+=1
self.__size__-=1
return True
def size(self):
return self.__size__
def leaders(self):
return self.__leaders__