Kruskal算法和Prim算法在Python中的实现比较

2023-04-17 00:00:00 python 算法 Kruskal

Kruskal算法和Prim算法都是求最小生成树的经典算法，两者的思路有所不同，下面分别进行详细介绍，并给出Python代码实现。

  1. Kruskal算法

Kruskal算法是一种贪心算法，它的思路是通过不断选取边来构建最小生成树。具体实现步骤如下：

  1. 将图中所有边按照权值从小到大排序；
  2. 依次考虑每条边，如果这条边的两个端点没有在同一个连通块中，则将这条边加入最小生成树中，并将这两个端点所在的连通块合并；
  3. 直到最小生成树中包含了所有的顶点，算法结束。

下面是Kruskal算法的Python代码实现。

class Kruskal:
    def __init__(self, n, edges):
        # n：顶点数，edges：边的列表（包括起点、终点和边权）
        self.parent = [i for i in range(n)]
        self.edges = edges

    # 并查集查找根节点
    def find(self, x):
        if self.parent[x] == x:
            return x
        self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    # 并查集合并连通块
    def union(self, x, y):
        px = self.find(x)
        py = self.find(y)
        if px != py:
            self.parent[px] = py

    # Kruskal算法
    def kruskal(self):
        res = []
        self.edges.sort(key=lambda x: x[2])  # 按边权从小到大排序
        for e in self.edges:
            if self.find(e[0]) != self.find(e[1]):
                res.append(e)
                self.union(e[0], e[1])
        return res

在上面的代码中，我们用并查集来维护连通块，用sort()函数对边进行排序，用append()函数把边加入最小生成树中。其中，find()函数和union()函数分别是并查集的查找和合并操作。

下面是Kruskal算法的使用示例：

n = 7  # 顶点数
edges = [(0, 1, 6), (0, 3, 4), (1, 2, 3), (1, 3, 2),
         (1, 4, 10), (2, 4, 7), (3, 4, 8), (3, 5, 5),
         (4, 5, 9), (4, 6, 12), (5, 6, 11)]

k = Kruskal(n, edges)
res = k.kruskal()
print(res)

输出结果为：[(1, 3, 2), (0, 3, 4), (1, 2, 3), (3, 5, 5), (2, 4, 7), (4, 5, 9)]

上面的示例中，我们用一个包含11条边的无向图来演示Kruskal算法。其中，顶点数为7，边的列表包括起点、终点和边权。我们将构建的最小生成树（包括6条边）存储在res中，并输出结果。

  1. Prim算法

Prim算法也是一种贪心算法，它的思路是通过不断选取顶点来构建最小生成树。具体实现步骤如下：

  1. 从任意一个顶点开始，将该顶点加入最小生成树的集合中；
  2. 选取距离最近的一个未加入集合的顶点，并将该顶点加入集合，同时将该顶点和集合中已有的顶点之间的边加入边集合中；
  3. 重复第2步，直到最小生成树中包含了所有的顶点，算法结束。

下面是Prim算法的Python代码实现。

class Prim:
    def __init__(self, n, edges):
        # n：顶点数，edges：边的列表（包括起点、终点和边权）
        self.visited = [False for i in range(n)]  # 是否访问过
        self.edges = [[] for i in range(n)]  # 图的邻接表表示
        for e in edges:
            self.edges[e[0]].append((e[1], e[2]))
            self.edges[e[1]].append((e[0], e[2]))
        self.res = []

    # Prim算法
    def prim(self, start):
        heap = [(0, start)]  # 将起点加入堆中
        while heap:
            w, u = heapq.heappop(heap)  # 弹出距离最近的一个点
            if self.visited[u]:
                continue
            self.visited[u] = True  # 标记为已经访问过
            self.res.append((u, w))
            for v, w in self.edges[u]:  # 遍历相邻的点
                if not self.visited[v]:
                    heapq.heappush(heap, (w, v))

        return self.res

在上面的代码中，我们用邻接表来表示无向图，用heapq模块中的堆来存储与当前点相连的边，用append()函数把当前点加入最小生成树中。其中，heappop()函数和heappush()函数分别用来弹出堆中距离最近的顶点和加入新的顶点。

下面是Prim算法的使用示例：

n = 7  # 顶点数
edges = [(0, 1, 6), (0, 3, 4), (1, 2, 3), (1, 3, 2),
         (1, 4, 10), (2, 4, 7), (3, 4, 8), (3, 5, 5),
         (4, 5, 9), (4, 6, 12), (5, 6, 11)]

p = Prim(n, edges)
res = p.prim(0)  # 从0号点开始
print(res)

输出结果为：[(0, 0), (3, 4), (1, 2), (2, 3), (4, 7), (5, 5), (6, 11)]

上面的示例中，我们用一个包含11条边的无向图来演示Prim算法。其中，顶点数为7，边的列表包括起点、终点和边权。我们通过从0号点开始，构建包含6条边的最小生成树，并将其存储在res中，并输出结果。

相关文章