Prim算法实现 -- 结合优先级队列

什么是Prim算法？

class Prim2:
    """
    P算法
    最小生成树算法 MST Minimal Spanning Tree
    保证整个拓扑图的所有路径之和最小
    """
    def __init__(self, graph):
        n = len(graph)
        # 存放横切边
        self.min_heap = []
        # 类似于visited数组，记录节点是否在mst中
        self.inMst = [False] * n
        self.weightSum = 0
        # 三元组
        self.graph = graph

    def prim(self):
        # 随便从一个节点切分，不妨从节点0开始
        self.inMst[0] = True
        self.cut(0)

        while self.min_heap:
            edge = heapq.heappop(self.min_heap)
            to = edge[1]
            weight = edge[2]

            if self.inMst[to]:
                continue
            self.weightSum += weight
            self.inMst[to] = True
            # 新一轮切分
            self.cut(to)

        return self.weightSum

    def cut(self, s):
        """
        将s的横切边加入优先队列
        :param s:
        :return:
        """
        for edge in self.graph[s]:
            dst = edge[1]
            if self.inMst[dst]:
                continue
            heapq.heappush(self.min_heap, edge)

    def allConnected(self):
        return all(self.inMst)