Python 利用heapq实现prim算法求解最小生成树 MST 并用graphviz直观表示

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格式图形化报告...

此即为最小生成树。