1+ # Basic implementation of the Union Find data structure
2+ # Assume we have n nodes labeled from 0 to n - 1
3+ 4+ class UnionFind :
5+ def __init__ (self , n ):
6+ # every node is originally its own parent
7+ self .par = [i for i in range (n )]
8+ # self.par = list(range(n)) -- also valid
9+ 10+ # every node originally is in its own
11+ # component of size 1 - this changes during
12+ # the union operation
13+ self .rank = [1 ] * n
14+ 15+ def find (self , n ) -> int :
16+ '''
17+ Finds the parent node of n
18+ '''
19+ 20+ # can be optimized with path compression
21+ while n != self .par [n ]:
22+ n = self .par [n ]
23+ return n
24+ 25+ 26+ def union (self , n1 , n2 ) -> bool :
27+ '''
28+ Connects two nodes together if not
29+ already connected
30+ '''
31+ 32+ # find the parent of node 1 and 2
33+ p1 = self .find (n1 )
34+ p2 = self .find (n2 )
35+ 36+ # nodes are already connected
37+ # cannot union together
38+ if p1 == p2 :
39+ return False
40+ 41+ # for efficiency, make bigger component
42+ # parent of smaller component - reduces
43+ # number of steps we have to take in find()
44+ 45+ if self .rank [p1 ] >= self .rank [p2 ]:
46+ # p2 is smaller, so when union it has a
47+ # new parent, p1
48+ self .par [p2 ] = p1
49+ 50+ # p1 gets all the nodes of p2, increasing
51+ # its rank, or size
52+ self .rank [p1 ] += self .rank [p2 ]
53+ else :
54+ self .par [p1 ] = p2
55+ self .rank [p2 ] += self .rank [p1 ]
56+ 57+ return True
58+ 59+ def nodes_connected (self , n1 , n2 ) -> bool :
60+ '''
61+ Returns if two nodes are connected
62+ '''
63+ 64+ # connected if parent is the same
65+ return self .find (n1 ) == self .find (n2 )
66+ 67+ 68+ 69+ def verify ():
70+ n = 7
71+ u = UnionFind (n )
72+ 73+ # False, nodes not connected
74+ print (u .nodes_connected (0 , 1 ))
75+ 76+ # True, just connected 0 and 1
77+ u .union (0 , 1 )
78+ print (u .nodes_connected (0 , 1 ))
79+ 80+ # Rank is 2, includes 0 and 1
81+ print (u .rank [0 ])
82+ 83+ u .union (4 , 5 )
84+ u .union (1 , 4 )
85+ 86+ # True, 0 - 1 and 4 - 5 are connected
87+ # 1 to 4 connects both components
88+ print (u .nodes_connected (0 , 5 ))
0 commit comments