from heapq import heapify, heappop, heappush
from collections import defaultdict


def prim(nodes, edges):
    #######conn是邻接表，字典形式
    conn = defaultdict(list)
    ####建立邻接表，如{‘A‘:[(7,‘A‘,‘B‘),(5,‘A‘,‘D‘)]}
    for n1, n2, c in edges:
        conn[n1].append((c, n1, n2))
        conn[n2].append((c, n2, n1))

    ##### mst就是存边结果，[(‘A‘,‘D‘,5)……]
    mst = []
    ##### used就是存已经放进mst中的点，
    used = set(nodes[0])
    ##### usable_edges 为建立好的树（采用堆排序建立的树）
    usable_edges = conn[nodes[0]][:]
    #####堆排序，保证堆顶的边权重为最小。
    heapify(usable_edges)

    ##### 堆中弹出的点为最后加入的点，其连接的其他边如果不在used里，则加入堆中
    while usable_edges:
        cost, n1, n2 = heappop(usable_edges)
        if n2 not in used:
            used.add(n2)
            mst.append((n1, n2, cost))
            ####conn最后加入的点的其他边的也加入堆中
            for e in conn[n2]:
                if e[2] not in used:
                    heappush(usable_edges, e)
    return mst


nodes = list("ABCDEFG")
edges = [("A", "B", 7), ("A", "D", 5),
         ("B", "C", 8), ("B", "D", 9),
         ("B", "E", 7), ("C", "E", 5),
         ("D", "E", 15), ("D", "F", 6),
         ("E", "F", 8), ("E", "G", 9),
         ("F", "G", 11)]
print ("prim:", prim(nodes, edges))

from collections import defaultdict
from heapq import heapify, heappush, heappop
import sys


def prim(m, arr):
    if len(arr) == 0 or len(arr[0]) == 0:
        return 0
    nodes = [i for i in range(m)]
    graph = defaultdict(list)
    res = 0
    for i in range(len(arr)):
        for j in range(1,len(arr[i])):
            graph[i].append([arr[i][j], i, i + j + 1])

    use_node = set([nodes[0]])
    heap_node = graph[0]
    heapify(heap_node)
    while heap_node:
        val, u, v = heappop(heap_node)
        if v not in use_node:
            use_node.add(v)
            res += val
            for e in graph[v]:
                if e[2] not in use_node:
                    heappush(heap_node, e)
    return res


m = 3
arr = [[2, 1, 3], [2, 1], [3]]
# arr = [[1,2,2],[1,2],[1]]
print(prim(m, arr))