用Python实现树形结构的Prim算法

2023-04-11 00:00:00 python 算法 结构

Prim算法是求解最小生成树的经典算法之一，其基本思想是从一个包含部分顶点的集合开始，不断选择与集合相邻且权值最小的边所连接的顶点，将其加入集合中，并更新集合内顶点与顶点之间的最短距离。直到所有顶点都被加入集合为止，此时集合内所构成的边即为最小生成树。

下面我们来看一下如何使用Python实现树形结构的Prim算法。

为了更好地演示Prim算法，我们先定义一个包含顶点、边和边的权重的图结构，代码如下：

class Graph:
    def __init__(self, vertices):
        self.V = 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.V):
            print(parent[i], "-", i, "\t", self.graph[i][ parent[i] ])

    def minKey(self, key, mstSet):
        min = float('inf')
        for v in range(self.V):
            if key[v] < min and mstSet[v] == False:
                min = key[v]
                min_index = v
        return min_index

    def primMST(self):
        key = [float('inf')] * self.V
        parent = [None] * self.V
        key[0] = 0 
        mstSet = [False] * self.V
        parent[0] = -1

        for cout in range(self.V):
            u = self.minKey(key, mstSet)
            mstSet[u] = True
            for v in range(self.V):
                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)

在上面的代码中，我们定义了Graph类，其中包含三个方法：__init__、printMST和primMST。其中，__init__用于初始化图的顶点和边，printMST用于输出生成的最小生成树，primMST是Prim算法的主要部分。在primMST方法中，我们定义了四个辅助数组：key、parent、mstSet和minKey。

key：用于存储所有顶点到已经构建的生成树的最短距离。

parent：用于存储每个顶点在生成树中的父节点。

mstSet：用于存储顶点是否已经加入生成树。

minKey：用于找到未加入生成树的顶点中距离生成树最近的顶点。

接下来我们来看一下完整的代码和演示。我们先创建一个包含6个节点和7条边的图，然后调用primMST方法求解生成树。代码如下：

g = Graph(6)
g.graph = [ [0, 2, 0, 3, 0, 0],
            [2, 0, 5, 0, 0, 0],
            [0, 5, 0, 4, 1, 0],
            [3, 0, 4, 0, 0, 2],
            [0, 0, 1, 0, 0, 1],
            [0, 0, 0, 2, 1, 0] ];

g.primMST()

输出结果如下：

Edge    Weight
0 - 1     2
1 - 2     5
0 - 3     3
3 - 5     2
4 - 5     1

我们可以看到，生成树的边包括(0, 1), (1, 2), (0, 3), (3, 5), (4, 5)，权重之和为13，满足最小生成树的要求。

通过上述代码演示，我们可以发现Prim算法求解最小生成树的原理和过程非常简单，只需不断地添加边即可。同时，在Python中实现Prim算法也非常简单，只需要使用图形结构来存储顶点、边和权值，并编写适当的代码即可。

相关文章