题解 | 最小生成树

#
# 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
#
# 返回最小的花费代价使得这n户人家连接起来
# @param n int整型 n户人家的村庄
# @param m int整型 m条路
# @param cost int整型二维数组 一维3个参数，表示连接1个村庄到另外1个村庄的花费的代价
# @return int整型
#
class Solution:
    def miniSpanningTree(self, n: int, m: int, cost: List[List[int]]) -> int:
        # write code here
        # prim算法
        import heapq
        # 邻接表、访问记录、用于弹出最小值的堆
        # 构造邻接表方便查找可选边及权重
        graph = [[] for _ in range(n+1)]
        for u,v,w in cost:
            graph[u].append((v,w))
            graph[v].append((u,w))

        visited =[False]*(n+1)
        total_weight = 0
        visited_count = 0
        min_heap = [(0,1)]

        while min_heap and visited_count<n:
            w,node = heapq.heappop(min_heap)

            if visited[node]:
                continue

            visited[node] = True
            visited_count += 1
            total_weight += w

            for neighbor,w in graph[node]:
                if not visited[neighbor]:
                    heapq.heappush(min_heap,(w,neighbor))

        return total_weight