python--数据结构--并查集

# mf_set
from typing import List, Sequence


class TNode:
    def __init__(self, data, parent: int):
        self.data = data
        self.parent = parent


class SeqList:  # S
    def __init__(self, elem: List):
        self.elem = elem
        self.last = len(self.elem) - 1


class MFSet:  # SS
    def __init__(self):
        self.tree: List[TNode] = []
        self.nodenum: int = 0

    def initialization(self, s: SeqList):
        """
        用s中的n个元素构成n个单根树，代表n个单元素集合S0,S1,S2,...，Sn-1，这n个单根树构成一个森林，代表并查集SS
        :param s: n个元素列表
        :return: None
        """
        self.nodenum = s.last + 1
        for item in s.elem:
            self.tree.append(TNode(item, -1))

    def locate(self, x) -> int:

        for i in range(self.nodenum):
            if self.tree[i].data == x:
                return i
        return -1  # 表示在并查集中没有找到x

    def find_1(self, x):
        """
        确定x属于并查集tree[]中哪个子集合。如果不属于tree[]中任一子集，则返回-1，否则返回所在子集树的根节点下标
        :param x:
        :return:
        """
        pos = self.locate(x)
        if pos < 0:
            return -1  # 表示在并查集中没有找到x
        i = pos
        while self.tree[i].parent > 0:
            i = self.tree[i].parent
        return i

    def find_2(self, x):
        """
        确定x属于并查集tree[]中哪个子集合，同时调整子集树，降低其高度。如果x属于tree[]中任一子集，则返回-1，否则首先找到x
        所在子集树的根root，然后将x及x的所有则先（除了root）均改为root的子节点，最后返回root。
        :param x:
        :return:
        """
        pos = self.locate(x)  # 确定x在self.tree[]中的下标
        if pos < 0:  # 如果x不属于self.tree[]中的任一子集，则返回-1
            return -1  # 表示在并查集中没有找到x
        i = pos  # 从pos开始，沿着双亲指针查找根节点
        while self.tree[i].parent >= 0:
            i = self.tree[i].parent
        root = i  # 记录x所在子集树的根结点下标
        i = pos  # 从pos开始，将x及x的所有祖先（除了root）均改为root的子结点
        while i != root:
            temp = self.tree[i].parent
            self.tree[i].parent = root
            i = temp
        return root  # 返回x所在子集树的根结点下标

    def merge(self, root1: int, root2: int):
        """
        root1 和 root2 是并查集tree[]中两个互不相交的非空子集树的根，根结点的parent域存放树中结点数目的负值。本算法将
        结点数目较少的子集树并入结点数目较多的子集树
        :param root1: 非空子集树的根
        :param root2: 与root1互不相交的非空子集树的根
        :return: 成功返回"OK"，在并查集中没有找到root1或root2返回"ERROR"
        """
        if root1 < 0 or root1 > self.nodenum - 1:
            return "ERROR"  # 表示在并查集中没有找到root1
        if root2 < 0 or root2 > self.nodenum - 1:
            return "ERROR"  # 表示在并查集中没有找到root2
        if self.tree[root1].parent < self.tree[root2].parent:  # 第一棵子集树中结点数目较多
            self.tree[root1].parent += self.tree[root2].parent
            self.tree[root2].parent = root1
        else:  # 第二棵子集树中结点数目较多
            self.tree[root2].parent += self.tree[root1].parent
            self.tree[root1].parent = root2
        return "OK"

# test_mf_set
from mf_set import MFSet, SeqList


seq_list = ['A', 'B', 'C']
mf_set = MFSet()
mf_set.initialization(SeqList(seq_list))

a_i = mf_set.find_2('B')
b_i = mf_set.find_2('C')
print(a_i, b_i)
mf_set.merge(a_i, b_i)

a_i = mf_set.find_2('A')
b_i = mf_set.find_2('C')
print(a_i, b_i)
mf_set.merge(a_i, b_i)
for item in mf_set.tree:
    print((item.data, item.parent))


"""
运行结果：
1 2
0 2
('A', 2)
('B', 2)
('C', -3)

Process finished with exit code 0
"""