python實現並查集

並查集是這樣的數據結構：有一大堆的數據，把一些元素放在一個集合當中，另外一些元素放在另一個一個集合當中。

對於它的操作有：查看兩個元素是否在一個集合當中、合並兩個元素。 合並的時候采取的策略是這樣的：將兩個元素所在的集合的所有元素一起放入一個集合當中。

這里使用兩個字典來實現並查集：一個字典保存當前節點的父節點的信息，另外一個保持父節點大小的信息。

class UnionFindSet(object):
    """並查集"""
    def __init__(self, data_list):
        """初始化兩個字典，一個保存節點的父節點，另外一個保存父節點的大小
        初始化的時候，將節點的父節點設為自身，size設為1"""
        self.father_dict = {}
        self.size_dict = {}

        for node in data_list:
            self.father_dict[node] = node
            self.size_dict[node] = 1

    def find_head(self, node):
        """使用遞歸的方式來查找父節點

        在查找父節點的時候，順便把當前節點移動到父節點上面
        這個操作算是一個優化
        """
        father = self.father_dict[node]
        if(node != father):
            father = self.find_head(father)
        self.father_dict[node] = father
        return father

    def is_same_set(self, node_a, node_b):
        """查看兩個節點是不是在一個集合里面"""
        return self.find_head(node_a) == self.find_head(node_b)

    def union(self, node_a, node_b):
        """將兩個集合合並在一起"""
        if node_a is None or node_b is None:
            return

        a_head = self.find_head(node_a)
        b_head = self.find_head(node_b)

        if(a_head != b_head):
            a_set_size = self.size_dict[a_head]
            b_set_size = self.size_dict[b_head]
            if(a_set_size >= b_set_size):
                self.father_dict[b_head] = a_head
                self.size_dict[a_head] = a_set_size + b_set_size
            else:
                self.father_dict[a_head] = b_head
                self.size_dict[b_head] = a_set_size + b_set_size

if __name__ == '__main__':
    a = [1,2,3,4,5]
    union_find_set = UnionFindSet(a)
    union_find_set.union(1,2)
    union_find_set.union(3,5)
    union_find_set.union(3,1)
    print(union_find_set.is_same_set(2,5))  # True