629. Minimum Spanning Tree - Prim

way 2. Min heap, Prim 方法, 从某点出发. O(V log V)
思路:
从每一个顶点开始, 把它的所有边加进 min heap. heap node: (weight, city1, city2, vertex),
循环, while h:
- 弹出最小边的另一个顶点,
- 跳过 used 顶点,
- 最短边加进结果集,
 -- 子循环, 把此顶点的所有边加进 min heap, 跳过 used 顶点,
- 直到加满树的边数, if len(res) == len(V) - 1: 边数等于顶点数减一, 是一棵树.

Note, 此题需要保留 city1, city2 的顺序, 并且 heap 按照 (cost, city1, city2) 排序,

'''
Definition for a Connection
class Connection:

    def __init__(self, city1, city2, cost):
        self.city1, self.city2, self.cost = city1, city2, cost
'''
import heapq, collections
class Solution:
    def lowestCost(self, C):
        # init
        graph = collections.defaultdict(list)
        for c in C:
            graph[c.city1].append((c.city2, c))
            graph[c.city2].append((c.city1, c))
        h, res = [], []
        n = len(graph)
        # used = set([C[0].city1])
        used = set()
        heapq.heappush(h, (-1, "", "", C[0].city1))

        while h: # and len(res) < n - 1:
            cost, c1, c2, u = heapq.heappop(h)
            if u in used:
                continue
            if c1 != "":
                res.append(Connection(c1, c2, cost))
            used.add(u)
            # print(u, c1, c2, cost)
            for v, c in graph[u]:
                if v in used:
                    continue
                heapq.heappush(h, (c.cost, c.city1, c.city2, v))
        if len(res) != n - 1:
            return []
        res.sort(key = lambda x: (x.cost, x.city1, x.city2))
        return res