heapq实现了小顶堆,主要用在快速选出最小边。并查集用在快速判断是否成环。
代码实现1
#!/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格式图形化报告...
此即为最小生成树。