2020数据结构小学期（四）——Prim算法

4、Prim算法

输入：无向图（顶点序列，边序列）

功能要求：输出最小生成树的各组成边及最小生成树的权值

源码：

import numpy as np


class MinTree(object):
    def create_graph(self, graph, vertex: int, data: [], weight: []):
        for i in range(vertex):
            graph.data[i] = data[i]
            for j in range(vertex):
                graph.weight[i][j] = weight[i][j]

    def show_graph(self, graph):
        for link in graph.weight:
            print(link)

    def prim_algorithm(self, graph, v: int):
        total = 0
        visited = [0] * graph.vertex
        visited[v] = 1
        h1 = -1
        h2 = -1
        min_weight = 10000
        for k in range(1, graph.vertex):
            for i in range(graph.vertex):
                for j in range(graph.vertex):
                    if visited[i] == 1 and visited[j] == 0 and graph.weight[i][j] < min_weight:
                        min_weight = graph.weight[i][j]
                        h1 = i
                        h2 = j
            print('边 %s -> %s 权值： %d ' % (graph.data[h1], graph.data[h2], min_weight))
            total += min_weight
            visited[h2] = 1
            min_weight = 10000
        print('最小生成树的权值： ', total)


class MGraph(object):
    def __init__(self, vertex):
        self.vertex = vertex
        self.data = vertex * [0]
        self.weight = [[0 for row in range(vertex)] for col in range(vertex)]


if __name__ == '__main__':
    """
    vertex_data = ['A', 'B', 'C', 'D', 'E', 'F', 'G']
    weights = [[10000, 5, 7, 10000, 10000, 10000, 2],
               [5, 10000, 10000, 9, 10000, 10000, 3],
               [7, 10000, 10000, 10000, 8, 10000, 10000],
               [10000, 9, 10000, 10000, 10000, 4, 10000],
               [10000, 10000, 8, 10000, 10000, 5, 4],
               [10000, 10000, 10000, 4, 5, 10000, 6],
               [2, 3, 10000, 10000, 4, 6, 10000]]
    """

    vertex_str = input("请输入顶点序列（以空格分隔）：")
    vertex_list = vertex_str.split(" ")
    vertex_data = [char for char in vertex_list]
    weight_str = input("请输入边序列（以空格分隔）：")
    weight_list = weight_str.split(" ")
    weight_list = list(map(int, weight_list))
    weight_list = np.array(weight_list)
    weights = weight_list.reshape(len(vertex_data), len(vertex_data))
    g = MGraph(len(vertex_data))
    tree = MinTree()
    tree.create_graph(g, len(vertex_data), vertex_data, weights)
    tree.show_graph(g)
    tree.prim_algorithm(g, 0)