python 基于字典的并查集模板 + Kruskal算法模板

Kruskal算法

Kruskal基于并查集算法来找到最小生成树。

算法步骤

1、初始化并查集。计算权重
2、根据权重对边进行排序。
3、遍历排序后的边，判断该边对应的两个顶点是否联通，然后将其联通，并加入生成树

按照模板简单来说
先初始化并查集
将节点add进并查集
计算边的权重并加上边放在新的列表edges中
在edges中对边按照权重进行排序

最后遍历并查集：

判断是否连通：
若连通则跳过
若不连通则连通两节点
并将结果加上权重

Code

模板中 points 举例：

points = [[0,0],[2,2],[3,10],[5,2],[7,0]]

模板（以 points 为输入数组）：

class UnionFind:
    def __init__(self):
        """
        记录每个节点的父节点
        """
        self.father = {}

    def find(self, x):
        """
        查找根节点
        路径压缩
        """
        root = x

        while self.father[root] is not None:
            root = self.father[root]

        # 路径压缩
        while x != root:
            original_father = self.father[x]
            self.father[x] = root
            x = original_father

        return root

    def merge(self, x, y):
        """
        合并两个节点
        """
        root_x, root_y = self.find(x), self.find(y)

        if root_x != root_y:
            self.father[root_x] = root_y

    def is_connected(self, x, y):
        """
        判断两节点是否相连
        """
        return self.find(x) == self.find(y)

    def add(self, x):
        """
        添加新节点
        """
        if x not in self.father:
            self.father[x] = None

    # KRUSKAL算法实现


def kruskal(points):
    if not points:
        return 0

    n = len(points)

    # 根据Kruskal算法三部曲：
    # 1、构造权重边：因为Kruskal算法以边为维度，因此需要首先构造所有点两两相连时每条边的权重，在这里就是边的长度。

    edges = []
    # 这里需要两重循环，但是需要注意，第二层循环需要以第一层循环为起点（两两相连嘛）

    for node1 in range(n):
        for node2 in range(node1, n):
            p1 = points[node1]
            p2 = points[node2]
            # 这里需要计算权重，也就是边的长度
            edge = (node1, node2, abs(p1[0] - p2[0]) + abs(p1[1] - p2[1]))
            edges.append(edge)

    # 2、排序：对所有的边按照权重从小到大排序（题目要求是最小的花费）
    edges.sort(key=lambda item: item[-1])

    # 3、生成最小生成树：Kruskal是以并查集算法为基础的，因此这里就是通过并查集来构造联通分量。
    uf = UnionFind()

    for i in range(n):
        uf.add(i)

    res = 0
    # 因为在第二部的时候，已经将节点按照题目要求排过序了，因此这里直接使用并查集构造联通分量即可
    for node1, node2, dest in edges:
        if uf.is_connected(node1, node2):
            # 如果两个点已经连接了，跳过
            continue
        # 将花费加入结果中
        res += dest
        # 将两个点连接起来
        uf.merge(node1, node2)

    return res

具体实际调用举例：

class UnionFind:
    def __init__(self):
        """
        记录每个节点的父节点
        """
        self.father = {}

    def find(self, x):
        """
        查找根节点
        路径压缩
        """
        root = x

        while self.father[root] is not None:
            root = self.father[root]

        # 路径压缩
        while x != root:
            original_father = self.father[x]
            self.father[x] = root
            x = original_father

        return root

    def merge(self, x, y):
        """
        合并两个节点
        """
        root_x, root_y = self.find(x), self.find(y)

        if root_x != root_y:
            self.father[root_x] = root_y

    def is_connected(self, x, y):
        """
        判断两节点是否相连
        """
        return self.find(x) == self.find(y)

    def add(self, x):
        """
        添加新节点
        """
        if x not in self.father:
            self.father[x] = None

    # KRUSKAL算法实现


def kruskal(points):
    if not points:
        return 0

    n = len(points)

    # 根据Kruskal算法三部曲：
    # 1、构造权重边：因为Kruskal算法以边为维度，因此需要首先构造所有点两两相连时每条边的权重，在这里就是边的长度。

    edges = []
    # 这里需要两重循环，但是需要注意，第二层循环需要以第一层循环为起点（两两相连嘛）

    for node1 in range(n):
        for node2 in range(node1, n):
            p1 = points[node1]
            p2 = points[node2]
            # 这里需要计算权重，也就是边的长度
            edge = (node1, node2, abs(p1[0] - p2[0]) + abs(p1[1] - p2[1]))
            edges.append(edge)

    # 2、排序：对所有的边按照权重从小到大排序（题目要求是最小的花费）
    edges.sort(key=lambda item: item[-1])

    # 3、生成最小生成树：Kruskal是以并查集算法为基础的，因此这里就是通过并查集来构造联通分量。
    uf = UnionFind()

    for i in range(n):
        uf.add(i)

    res = 0
    # 因为在第二部的时候，已经将节点按照题目要求排过序了，因此这里直接使用并查集构造联通分量即可
    for node1, node2, dest in edges:
        if uf.is_connected(node1, node2):
            # 如果两个点已经连接了，跳过
            continue
        # 将花费加入结果中
        res += dest
        # 将两个点连接起来
        uf.merge(node1, node2)

    return res


points = [[0, 0], [2, 2], [3, 10], [5, 2], [7, 0]]
print(kruskal(points))