Python中使用Prim算法求最小生成树

2023-04-11 00:00:00 算法 生成 最小

Prim算法是一种用于计算无向加权图的最小生成树(MST)的算法。它在一个连通的加权无向图中选取一个顶点作为起点,然后通过不断扩展树的边缘来构建最小生成树。在这个过程中,我们必须维护两个集合,一个是已选择的边构成的集合,另一个是未选择的边构成的集合。

下面是使用Python实现Prim算法的例子:

import sys

class Graph(object):
    def __init__(self, vertices):
        self.vertices = vertices
        self.graph = [[0 for column in range(vertices)] for row in range(vertices)]

    def printMST(self, parent):
        print("Edge \tWeight")
        for i in range(1, self.vertices):
            print(parent[i], "-", i, "\t", self.graph[i][ parent[i] ])

    def minKey(self, key, mstSet):
        min = sys.maxsize

        for v in range(self.vertices):
            if key[v] < min and mstSet[v] == False:
                min = key[v]
                min_index = v

        return min_index

    def primMST(self):
        key = [sys.maxsize] * self.vertices
        parent = [None] * self.vertices # Array to store constructed MST
        key[0] = 0
        mstSet = [False] * self.vertices

        parent[0] = -1  # First node is always the root of

        for cout in range(self.vertices):

            u = self.minKey(key, mstSet)

            mstSet[u] = True

            for v in range(self.vertices):

                if self.graph[u][v] > 0 and mstSet[v] == False and key[v] > self.graph[u][v]:

                        key[v] = self.graph[u][v]
                        parent[v] = u

        self.printMST(parent)

g = Graph(5)
g.graph = [[0, 2, 0, 6, 0],
            [2, 0, 3, 8, 5],
            [0, 3, 0, 0, 7],
            [6, 8, 0, 0, 9],
            [0, 5, 7, 9, 0],
           ]

g.primMST()

上面的代码首先定义了一个Graph类,初始化它的节点数量和一个矩阵graph来表示边的权重。然后,我们实现了primMST函数来实现Prim算法。该算法的基本思想是从一个起点开始,每次找到当前已选择的边和未选择的边之间的最小权重边。将这条边加入已选择的边集合中,并将它的端点加入到节点集合中。在不断的执行这个过程直到生成的树的边的数量为节点数减1则停止。

上面的例子中,我们设置graph矩阵来代表一个边权重的无向加权图。然后,我们创建一个Graph对象,调用primMST函数,该函数根据已选边和未选边集合来计算最小生成树。在运行程序后,代码将输出最小生成树的边及其权重。

相关文章