For union by rank (without path optimization) (https://leetcode.com/explore/learn/card/graph/618/disjoint-set/3879/). If you look at the output in the test code, I think rank isn't being properly represented. I think the rank optimization assumes that there will also be the path optimization, but it is misleading to talk about them seperately.
# UnionFind class
class UnionFind:
def __init__(self, size):
self.root = [i for i in range(size)]
self.rank = [1] * size
def find(self, x):
while x != self.root[x]:
x = self.root[x]
return x
def union(self, x, y):
rootX = self.find(x)
rootY = self.find(y)
if rootX != rootY:
if self.rank[rootX] > self.rank[rootY]:
self.root[rootY] = rootX
elif self.rank[rootX] < self.rank[rootY]:
self.root[rootX] = rootY
else:
self.root[rootY] = rootX
self.rank[rootX] += 1
def connected(self, x, y):
return self.find(x) == self.find(y)
def getRank(self, x):
return self.rank[x]
# Test Case
uf = UnionFind(10)
# 1-2-5-6-7 3-8-9-4
print(uf.getRank(1)) # 1, good
uf.union(1, 2)
print(uf.getRank(1)) # 2, good
uf.union(2, 5)
uf.union(5, 6)
uf.union(6, 7)
print(uf.getRank(1)) # 2, bad should be 4
uf.union(3, 8)
uf.union(8, 9)
uf.union(9, 4)
# 1-2-5-6-7-3-8-9-4
uf.union(7, 3)
print(uf.getRank(1)) # 3, bad should be 9