Prim算法 Python实现代码

本文只对Prim算法进行实现，若需了解算法原理可参考文末链接。

import numpy as np
import random


# 随机生成有权矩阵
def generate_map(ver_num=6):
    graph = []
    for i in range(ver_num):
        raw = []
        for j in range(ver_num):
            if i == j:
                raw.append(np.inf)
            else:
                raw.append(random.randint(1, 31))
        graph.append(raw)
    print(graph)
    return graph


# 求已经确定的顶点集合与未选顶点集合中的最小边
def min_edge(select, candidate, graph):
    # 记录最小边权重
    min_weight = np.inf
    # 记录最小边
    v, u = 0, 0
    # 循环扫描已选顶点与未选顶点，寻找最小边
    for i in select:
        for j in candidate:
            # 如果存在比当前的最小边权重还小的边，则记录
            if min_weight > graph[i][j]:
                min_weight = graph[i][j]
                v, u = i, j
    # 返回记录的最小边的两个顶点
    return v, u


# prim算法
def prim(graph):
    # 顶点个数
    vertex_num = len(graph)
    # 存储已选顶点，初始化时可随机选择一个起点
    select = [0]
    # 存储未选顶点
    candidate = list(range(1, vertex_num))
    # 存储每次搜索到的最小生成树的边
    edge = []
    # 由于连接n个顶点需要n-1条边，故进行n-1次循环，以找到足够的边
    for i in range(1, vertex_num):
        # 调用函数寻找当前最小边
        v, u = min_edge(select, candidate, graph)
        # 添加到最小生成树边的集合中
        edge.append((v, u))
        # v是select中的顶点，u为candidate中的顶点，故将u加入candidate，以代表已经选择该顶点
        select.append(u)
        # 同时将u从candidate中删除
        candidate.remove(u)
    print(edge)
    return edge


def main():
    # graph = generate_map()
    graph = [[np.inf, 3, np.inf, np.inf, 6, 5],
                  [3, np.inf, 1, np.inf, np.inf, 4],
                  [np.inf, 1, np.inf, 6, np.inf, 4],
                  [np.inf, np.inf, 6, np.inf, 8, 5],
                  [6, np.inf, np.inf, 8, np.inf, 2],
                  [5, 4, 4, 5, 2, np.inf]]
    prim(graph)

Prim算法原理