algorithm-in-python/dataStructure/unionFindSet.py

19 lines
495 B
Python
Raw Normal View History

2020-04-15 12:28:20 +08:00
class unionFindSet:
def __init__(self, S):
self.S = {i: i for i in S}
2020-05-13 10:45:51 +08:00
self.size = {i: 1 for i in S}
2020-04-15 12:28:20 +08:00
def find(self, x):
if x != self.S[x]:
self.S[x] = self.find(self.S[x])
return self.S[x]
def union(self, a, b, key=lambda x: x):
x, y = sorted((self.find(a), self.find(b)), key=key)
2020-05-13 10:45:51 +08:00
self.S[y] = x
if x != y:
self.size[x] += self.size[y]
def getSize(self, x):
return self.size[self.find(x)]