并查集_python实现(模板)

关于并查集的概念和效率，这里不详述，网上一大堆，这里主要记录一下并查集的 python代码实现

三步走：
1.初始化pre和rangks数组，pre为每个元素的父结点（上一级），ranks为每个元素作为根节点时的树的秩（树的深度），一开始设置pre的每个值为每个元素自己，设置ranks每个值为0.
2.find() 查询该元素的首级，顺便做路径压缩
3.join() 合并两个集合，按秩合并，秩小的树的根节点指向秩大的树的根节点。

# 这里记录一下一种新的数据结构：并查集


def find(x, pre):
    """
    查找x的最上级（首级）
    :param x: 要查找的数
    :param pre: 每个元素的首级
    :return: 元素的上一级
    """
    root, p = x, x  # root:根节点， p:指针
    
    # 找根节点
    while root != pres[root]:
        root = pres[root]

    # 路径压缩，把每个经过的结点的上一级设为root（直接设为首级）
    while p != pres[p]:
        p, pres[p] = pres[p], root
    return root


def join(x, y, pre, ranks):
    """
    合并两个元素（合并两个集合）
    :param x: 第一个元素
    :param y: 第二个元素
    :param pre: 每个元素的上一级
    :param ranks: 每个元素作为根节点时的秩（树的深度）
    :return: None
    """
    h1, h2 = find(x, pre), find(y, pre)
    # 当两个元素不是同一组的时候才合并
    # 按秩合并
    if h1 != h2:
        if ranks[h1] < ranks[h2]:
            pre[h1] = h2
        else:
            pre[h2] = h1
            if ranks[h1] == ranks[h2]:
                ranks[h1] += 1


# 结点数
n = 10

# 边数据
data = [[0, 9], [9, 3], [1, 2], [2, 8], [4, 5], [6, 7], [0, 5], [6, 8]]

# pre一开始设置每个元素的上一级是自己，ranks一开始设置每个元素的秩为0
pre, ranks = [i for i in range(n)], [0] * n

for edge in data:
    join(edge[0], edge[1], pre, ranks)

print('pre:\t', pre)
print('ranks:\t', ranks)
print('idx:\t', list(range(n)))

# pre:	 [0, 6, 1, 0, 0, 4, 6, 6, 1, 0]
# ranks:	 [2, 1, 0, 0, 1, 0, 2, 0, 0, 0]
# idx:	 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]