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

heapq实现了小顶堆，主要用在快速选出最小边。并查集用在快速判断是否成环。

代码实现¹

#!/usr/bin/env python3
#kruskal.py
#fumiama 20201029

import sys, time
from heapq import heappop, heappush
from graphviz import Digraph

class UnionFindSet(object):
    def __init__(self, nodes):
        self.fatherMap = {}
        self.setNumMap = {}
        for node in nodes:
            self.fatherMap[node] = node
            self.setNumMap[node] = 1
    def findFather(self, node):
        father = self.fatherMap[node]
        if (node != father):
            father = self.findFather(father)
        self.fatherMap[node] = father
        return father
    def isSameSet(self, a, b):
        return self.findFather(a) == self.findFather(b)
    def union(self, a, b):
        if a is None or b is None: return
        aFather = self.findFather(a)
        bFather = self.findFather(b)
        if (aFather != bFather):
            aNum = self.setNumMap[aFather]
            bNum = self.setNumMap[bFather]
            if (aNum <= bNum):
                self.fatherMap[aFather] = bFather
                self.setNumMap[bFather] = aNum + bNum
                self.setNumMap.pop(aFather)
            else:
                self.fatherMap[bFather] = aFather
                self.setNumMap[aFather] = aNum + bNum
                self.setNumMap.pop(bFather)

dot = Digraph(comment='Gragh2Print')
dot.edge_attr.update(arrowhead='none')
dot.graph_attr['rankdir'] = 'LR'

edgeLinks = dict()
edgeWeightsOf = dict()
edgeWeights = []
size = 0
nodes = set()

def exitWithError(*error):
    print(*error)
    exit()

def printGraph2Pdf(grf): grf.render('graph-output/output_kruskal.gv', view=True)

def addEdge(a, b, w):
    global edgeLinks, edgeWeights, edgeWeightsOf
    if a not in edgeLinks: edgeLinks[a] = set()
    if b not in edgeLinks: edgeLinks[b] = set()
    edgeLinks[a].add(b)
    edgeLinks[b].add(a)
    wi = int(w)
    if wi not in edgeWeightsOf.keys(): edgeWeightsOf[wi] = []
    if {a, b} not in edgeWeightsOf[wi]: edgeWeightsOf[wi].append({a, b})
    heappush(edgeWeights, wi)

def loadGraph(fileName):
    try: f = open(fileName, 'r')
    except: exitWithError("打开文件失败, 请检查文件名是否正确或程序是否有权限访问")
    global size, edgeLinks
    size, edgeCount = map(int, f.readline().split())
    print("节点:", size, "边数:", edgeCount)
    re = f.readline().split()[0]   #从何处开始搜索
    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)
        nodes.add(a)
        nodes.add(b)
    f.close()
    return re

def signEdge(start, end):
    dot.edge(start, end, color = "red")

def kruskalMST(start):
    global edgeLinks, edgeWeights
    print("搜索起点:", start)
    #print("结点:", nodes)
    #print("边权堆:", edgeWeights)
    #print("边对点:", edgeWeightsOf)
    forest = UnionFindSet(nodes)
    while edgeWeights != []:
        minWeight = heappop(edgeWeights)
        #print("弹出:", minWeight)
        if minWeight in edgeWeightsOf.keys():
            minEdge = edgeWeightsOf[minWeight].pop()
            a = minEdge.pop()
            b = minEdge.pop()
            if edgeWeightsOf[minWeight] == []: edgeWeightsOf.pop(minWeight)
            if not forest.isSameSet(a, b):
                forest.union(a, b)
                signEdge(a, b)

if __name__ == '__main__':
    if len(sys.argv) != 2: exitWithError("用法:", sys.argv[0], "文件位置")
    else:
        time_start = time.time()
        kruskalMST(loadGraph(sys.argv[1]))
        print("用时:", time.time() - time_start)
        print("生成pdf格式图形化报告...")
        printGraph2Pdf(dot)

构造样例

16 26
2
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

检验程序

$ ./kruskal.py graph_edge_weight.txt 
节点: 16 边数: 26
搜索起点: 2
用时: 0.0009100437164306641
生成pdf格式图形化报告...

此即为最小生成树。

---

1. littlehaes python实现并查集 ↩︎