本文只對Prim算法進行實現(xiàn),若需了解算法原理可參考文末鏈接。
import numpy as np
import random
# 隨機生成有權(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)確定的頂點集合與未選頂點集合中的最小邊
def min_edge(select, candidate, graph):
# 記錄最小邊權(quán)重
min_weight = np.inf
# 記錄最小邊
v, u = 0, 0
# 循環(huán)掃描已選頂點與未選頂點,尋找最小邊
for i in select:
for j in candidate:
# 如果存在比當前的最小邊權(quán)重還小的邊,則記錄
if min_weight > graph[i][j]:
min_weight = graph[i][j]
v, u = i, j
# 返回記錄的最小邊的兩個頂點
return v, u
# prim算法
def prim(graph):
# 頂點個數(shù)
vertex_num = len(graph)
# 存儲已選頂點,初始化時可隨機選擇一個起點
select = [0]
# 存儲未選頂點
candidate = list(range(1, vertex_num))
# 存儲每次搜索到的最小生成樹的邊
edge = []
# 由于連接n個頂點需要n-1條邊,故進行n-1次循環(huán),以找到足夠的邊
for i in range(1, vertex_num):
# 調(diào)用函數(shù)尋找當前最小邊
v, u = min_edge(select, candidate, graph)
# 添加到最小生成樹邊的集合中
edge.append((v, u))
# v是select中的頂點,u為candidate中的頂點,故將u加入candidate,以代表已經(jīng)選擇該頂點
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)