本文只對(duì)Prim算法進(jìn)行實(shí)現(xiàn)，若需了解算法原理可參考文末鏈接。

import numpy as np
import random


# 隨機(jī)生成有權(quán)矩陣
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


# 求已經(jīng)確定的頂點(diǎn)集合與未選頂點(diǎn)集合中的最小邊
def min_edge(select, candidate, graph):
    # 記錄最小邊權(quán)重
    min_weight = np.inf
    # 記錄最小邊
    v, u = 0, 0
    # 循環(huán)掃描已選頂點(diǎn)與未選頂點(diǎn)，尋找最小邊
    for i in select:
        for j in candidate:
            # 如果存在比當(dāng)前的最小邊權(quán)重還小的邊，則記錄
            if min_weight > graph[i][j]:
                min_weight = graph[i][j]
                v, u = i, j
    # 返回記錄的最小邊的兩個(gè)頂點(diǎn)
    return v, u


# prim算法
def prim(graph):
    # 頂點(diǎn)個(gè)數(shù)
    vertex_num = len(graph)
    # 存儲(chǔ)已選頂點(diǎn)，初始化時(shí)可隨機(jī)選擇一個(gè)起點(diǎn)
    select = [0]
    # 存儲(chǔ)未選頂點(diǎn)
    candidate = list(range(1, vertex_num))
    # 存儲(chǔ)每次搜索到的最小生成樹(shù)的邊
    edge = []
    # 由于連接n個(gè)頂點(diǎn)需要n-1條邊，故進(jìn)行n-1次循環(huán)，以找到足夠的邊
    for i in range(1, vertex_num):
        # 調(diào)用函數(shù)尋找當(dāng)前最小邊
        v, u = min_edge(select, candidate, graph)
        # 添加到最小生成樹(shù)邊的集合中
        edge.append((v, u))
        # v是select中的頂點(diǎn)，u為candidate中的頂點(diǎn)，故將u加入candidate，以代表已經(jīng)選擇該頂點(diǎn)
        select.append(u)
        # 同時(shí)將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算法原理