并查集python实现

class UnionNode(object):
    def __init__(self,item):
        self.item = item
        self.parent = None
class SmallUnion:
    def __init__(self):
        self.root = None
        self.datas = []
        self.size = 0
    def add(self,item):
        self.size += 1
        newNode = UnionNode(item)
        if self.root == None:
            self.root = newNode
        else:
            newNode.parent = self.root
            self.datas.append(newNode)
    def union(self,other):
        assert(isinstance(other,SmallUnion))
        if other.root is None:
            return
        othroot = UnionNode(other.root.item)
        othdatas = []
        for i in other.datas:
            newNode = UnionNode(i.item)
            newNode.parent = othroot
            othdatas.append(newNode)
        othroot.parent = self.root
        for i in othdatas:
            self.patchPath(i)
    def patchPath(self,node):
        pNode,sNode = node,node
        while pNode.parent:
            pNode = pNode.parent
        sNode.parent = pNode
    def remove(self,item):
        if self.size == 0:
            return
        if self.size == 1:
            if item == self.root.item:
                self.root = None
                self.size -= 1
        if self.size > 1:
            if item == self.root.item:
                m = min([i.item for i in self.datas])
                self.root.item = m
                item = m
            foundIdx = -1
            for idx,node in enumerate(self.datas):
                if item == node.item:
                    foundIdx = idx
                    break
            if foundIdx != -1:
                self.datas.pop(foundIdx)
                self.size -= 1
    def __iter__(self):
        return UnionIterator(self.root,self.datas)
class UnionIterator:
    def __init__(self,root,datas):
        self.curIdx = 0
        self.root = root
        self.datas = datas
    def __iter__(self):
        return self
    def next(self):
        if self.curIdx > len(self.datas):
            raise StopIteration
        elif self.curIdx == 0:
            item = self.root.item
            self.curIdx += 1
            return item
        else:
            item = self.datas[self.curIdx - 1].item
            self.curIdx += 1
            return item

你可能感兴趣的:(并查集python实现)