Prim算法

普里姆算法查找最小生成树的过程，采用了贪心算法的思想。对于包含 N 个顶点的连通网，普里姆算法每次从连通网中找出一个权值最小的边，这样的操作重复 N-1 次，由 N-1 条权值最小的边组成的生成树就是最小生成树。 

from collections import defaultdict
from heapq import heapify, heappop, heappush
def prim(map,start):
    #最小生成树
    mst = []
    #存已连接点集合
    city = {start}
    #存待走边集合
    next = map[start]#
    #待走边排序
    heapify(next)#按距离建立小顶堆
    while next:
        #取出最小边
        distance, head, tail = heappop(next)
        #如果下个路口没走过
        if tail not in city:
            #加入已连接集合
            city.add(tail)
            #加入最小生成树
            mst.append((head, tail, distance))
            # 没走过的相邻边加入
            for e in map[tail]:
                if e[2] not in city:
                    heappush(next, e)
    if len(mst) < len(map) - 1: return -1 #mst长度小于n-1则不联通
    else :
        min = 0
        for i,j,x in mst:
            min +=x
        return min
n,m = (int(i) for i in input().split(" "))
map = defaultdict(list)
for _ in range(m):
    begin,end,distance = (int(i) for i in input().split(" "))
    map[begin].append((distance, begin, end))
    map[end].append((distance, end, begin))
print(prim(map,1))

'''
7 11
1 2 7
2 3 8
1 4 5
4 2 9
2 5 7
3 5 5
4 5 15
4 6 6
6 5 8
5 7 9
6 7 11
'''