并查集 Python实现

# 并查集实现
class Node:
    pass

class UnionFindSet:
    def __init__(self, nodes):
        self.fatherDict = dict()
        self.sizeDict = dict()
        for node in nodes:
            self.fatherDict[node] = node
            self.sizeDict[node] = 1

    # def findHead(self, node):
    #     father = self.fatherDict.get(node)
    #     if node!=father:
    #         father = self.findHead(father)
    #     self.fatherDict[node] = father
    #     return father

    def findHead(self, node):           # 找到当前节点的头
        stack = []                      # 每次查询都会优化，经过的节点会直接指向头结点
        father = self.fatherDict[node]
        while node != father:
            stack.append(node)
            node = father
            father = self.fatherDict[node]
        while len(stack) > 0:
            self.fatherDict[stack.pop()] = father
        return father

    def isSameSet(self, a, b):
        return self.findHead(a) == self.findHead(b)

    def uion(self, a, b):                  # 两集合合并
        if a is None or b is None:
            return
        aHead = self.findHead(a)
        bHead = self.findHead(b)
        if aHead != bHead:
            aSize = self.sizeDict[aHead]
            bSize = self.sizeDict[bHead]
            if aSize <= bSize:
                self.fatherDict[aHead] = bHead
                self.sizeDict[bHead] = aSize + bSize
            else:
                self.fatherDict[bHead] = aHead
                self.sizeDict[aHead] = aSize + bSize