heapq实现了小顶堆，主要用在快速选出最小边
代码实现
#!/usr/bin/env python3
#prim.py
#fumiama 20201027
import sys, time
from heapq import heapify, heappop, heappush
from graphviz import Digraph
time_start = time.time()
#创建图
dot = Digraph(comment='Gragh2Print')
dot.edge_attr.update(arrowhead='none')
dot.graph_attr['rankdir'] = 'LR'
edgeLinks = dict() #点到点的映射关系
edgeWeights = dict() #边的权重
size = 0 #总的点的个数
def exitWithError(*error):
print(*error)
exit()
def printGraph2Pdf(grf): grf.render('graph-output/output.gv', view=True)
def addEdge(a, b, w):
global edgeLinks, edgeWeights
if a not in edgeLinks: edgeLinks[a] = set()
if b not in edgeLinks: edgeLinks[b] = set()
edgeLinks[a].add(b) #无向图添加双向边
edgeLinks[b].add(a)
if a not in edgeWeights.keys(): edgeWeights[a] = dict()
if b not in edgeWeights.keys(): edgeWeights[b] = dict()
edgeWeights[a][b] = edgeWeights[b][a] = w #无向图双向权重相等
def loadGraph(fileName):
try: f = open(fileName, 'r')
except: exitWithError("打开文件失败, 请检查文件名是否正确或程序是否有权限访问")
global size, edgeLinks
size, edgeCount = map(int, f.readline().split())
print("节点:", size, "边数:", edgeCount)
for i in range(1, size+1): dot.node(str(i), str(i))
for i in range(edgeCount):
a, b, w = f.readline().split()
addEdge(a, b, w)
dot.edge(a, b, w)
re = f.readline() #从何处开始搜索
f.close()
return re
def signEdge(start, end): #标红路径
dot.edge(start, end, color = "red")
def primMST(start):
global edgeLinks, edgeWeights
weightFromNode = dict() #从某一个点出发的边的权重到目标点的映射
print("起点:", start)
v = [start] #已搜索点集合
heap = [] #优先队列
while len(v) < size:
print("v:", v)
weightFromNode[v[-1]] = dict()
for destination in edgeLinks[v[-1]]:
evd = edgeWeights[v[-1]][destination]
evdi = int(edgeWeights[v[-1]][destination])
if destination in v: #避免成环
if destination in weightFromNode and evd in weightFromNode[destination].keys() and v[-1] in weightFromNode[destination][evd]:
weightFromNode[destination][evd].remove(v[-1])
if weightFromNode[destination][evd] == set(): weightFromNode[destination].pop(evd)
if evdi in heap:
heap.remove(evdi)
heapify(heap)
else:
if evd not in weightFromNode[v[-1]]: weightFromNode[v[-1]][evd] = set()
if destination not in weightFromNode[v[-1]][evd]:
heappush(heap, evdi)
weightFromNode[v[-1]][evd].add(destination)
#print("heap:", heap)
#print("weightFromNode:", weightFromNode)
shortEdgeW = str(heappop(heap))
#print("pop:", shortEdgeW)
stack = []
found = False
for start in weightFromNode.keys():
if shortEdgeW in weightFromNode[start]:
dest = weightFromNode[start][shortEdgeW].pop()
if weightFromNode[start][shortEdgeW] == set(): weightFromNode[start].pop(shortEdgeW)
if dest not in v:
signEdge(start, dest)
v.append(dest)
found = True
elif weightFromNode[start] == {
}: stack.append(start)
if found: break
for key in stack: weightFromNode.pop(key)
if __name__ == '__main__':
if len(sys.argv) != 2: exitWithError("用法:", sys.argv[0], "文件位置")
else:
primMST(loadGraph(sys.argv[1]))
print("用时:", time.time() - time_start)
print("生成pdf格式图形化报告...")
printGraph2Pdf(dot)
构造样例
16 26
1 2 5
1 5 6
2 3 2
2 4 5
3 4 9
3 5 12
4 5 4
5 7 8
5 14 3
5 6 6
6 9 5
7 4 9
9 4 23
10 6 14
11 5 56
12 7 87
13 8 4
13 6 7
15 13 8
16 10 2
8 15 1
13 10 1
9 13 2
15 16 1
8 12 10
11 16 7
2
检验程序
$ ./prim.py graph_edge_weight.txt
节点: 16 边数: 26
起点: 2
v: ['2']
v: ['2', '3']
v: ['2', '3', '1']
v: ['2', '3', '1', '4']
v: ['2', '3', '1', '4', '5']
v: ['2', '3', '1', '4', '5', '14']
v: ['2', '3', '1', '4', '5', '14', '6']
v: ['2', '3', '1', '4', '5', '14', '6', '9']
v: ['2', '3', '1', '4', '5', '14', '6', '9', '13']
v: ['2', '3', '1', '4', '5', '14', '6', '9', '13', '10']
v: ['2', '3', '1', '4', '5', '14', '6', '9', '13', '10', '16']
v: ['2', '3', '1', '4', '5', '14', '6', '9', '13', '10', '16', '15']
v: ['2', '3', '1', '4', '5', '14', '6', '9', '13', '10', '16', '15', '8']
v: ['2', '3', '1', '4', '5', '14', '6', '9', '13', '10', '16', '15', '8', '11']
v: ['2', '3', '1', '4', '5', '14', '6', '9', '13', '10', '16', '15', '8', '11', '7']
用时: 0.0009667873382568359
生成pdf格式图形化报告...
此即为最小生成树。