591. 连接图 III

class ConnectingGraph3:
    """
    @param a: An integer
    @param b: An integer
    @return: nothing
    """
    def __init__(self, n):
        # initialize your data structure here.
        self.count = n
        self.father = {}
        for i in range(1, n + 1):
            self.father[i] = i 
            
    def connect(self, a, b):
        # write your code here
        root_a = self.find(a)
        root_b = self.find(b)

        if (root_a != root_b):
            #a的父节点指向b（均是根节点,根节点连接）
            self.father[root_a] = root_b 
            self.count -= 1 

    """
    @return: An integer 
    """
    def query(self):
        # write your code here
        return self.count

    #找根节点
    def find(self,num):
        origin_num = num
        #如果是相等的话，则直接返回本身
        if self.father[num] == num:
            return num

        #如果不相等的话，则找父节点,一直更新到根节点
        while num != self.father[num]:
            num = self.father[num]

        #最后，压缩路径
        #一直到最后一个根节点为止，原始值的所有父节点都指向根节点
        while self.father[origin_num] != num:
            temp = self.father[origin_num]
            self.father[origin_num] = num
            origin_num = temp  
        
        #最终返回根节点
        return num

class UnionFind():
    def __init__(self, n):
        self.father = {}
        self.size = {}
        for i in range(1, n + 1):
            self.father[i] = i 
            self.size[i] = 1
    
    def find(self,x):
        origin_x = x 
        if (self.father[x] == x):
            return x 
        
        #一直不停的找，找到根节点,一直更新到最后，则相等，x最后为根节点
        while x != self.father[x]:
            x = self.father[x]

        #压缩路径，所有该节点的父节点均指向根节点
        while self.father[origin_x] != x:
            temp = self.father[origin_x]
            self.father[x] = x 
            origin_x = temp
        
        #最终返回根节点
        return x 

    def union(self,a,b):
        root_a = self.find(a)
        root_b = self.find(b)

        #当a和b的根节点不同的时候，此时a的根节点指向b的根节点，或者相反也可以
        if (root_a != root_b):
            #a的根节点指向b的根节点
            self.father[root_a] = root_b 
            self.size[b] = self.size[a] + self.size[b]
    
    def query(self, num):
        return self.size[num]
            
obj = UnionFind(5)
print(obj.find(1))
obj.union(1,5)
print(obj.find(1))
print(obj.query(5))

1
5
2