并查集 Python实现

Posted 2020-11-11 icekx

 1 # 并查集实现
 2 class Node:
 3     pass
 4 
 5 class UnionFindSet:
 6     def __init__(self, nodes):
 7         self.fatherDict = dict()
 8         self.sizeDict = dict()
 9         for node in nodes:
10             self.fatherDict[node] = node
11             self.sizeDict[node] = 1
12 
13     # def findHead(self, node):
14     #     father = self.fatherDict.get(node)
15     #     if node!=father:
16     #         father = self.findHead(father)
17     #     self.fatherDict[node] = father
18     #     return father
19 
20     def findHead(self, node):           # 找到当前节点的头
21         stack = []                      # 每次查询都会优化，经过的节点会直接指向头结点
22         father = self.fatherDict[node]
23         while node != father:
24             stack.append(node)
25             node = father
26             father = self.fatherDict[node]
27         while len(stack) > 0:
28             self.fatherDict[stack.pop()] = father
29         return father
30 
31     def isSameSet(self, a, b):
32         return self.findHead(a) == self.findHead(b)
33 
34     def uion(self, a, b):                  # 两集合合并
35         if a is None or b is None:
36             return
37         aHead = self.findHead(a)
38         bHead = self.findHead(b)
39         if aHead != bHead:
40             aSize = self.sizeDict[aHead]
41             bSize = self.sizeDict[bHead]
42             if aSize <= bSize:
43                 self.fatherDict[aHead] = bHead
44                 self.sizeDict[bHead] = aSize + bSize
45             else:
46                 self.fatherDict[bHead] = aHead
47                 self.sizeDict[aHead] = aSize + bSize