GO的prim

Go 的 heap 包： Go 标准库中的 container/heap 包提供了堆（优先队列）的实现。 heap.Interface 接口： 这个包定义了一个 heap.Interface 接口，任何类型只要实现了这个接口，就可以使用 heap 包的功能。 需要实现的方法： 要满足 heap.Interface，需要实现以下方法：

Len() int Less(i, j int) bool Swap(i, j int) Push(x interface{}) Pop() interface{}

// 为了方便存储，将 1 - n 改成 0 - n-1
func minimumCost(n int, connections [][]int) int {
    var (
        cost = 0
        edges = make([][]Edge, n)
        visited = make([]bool, n)
        h = &EdgeHeap{}
    )
    heap.Init(h)
    visited[0] = true // 从第 0 个开始
    for _, connect := range connections {
        edges[connect[0] - 1] = append(edges[connect[0] - 1], Edge{connect[1] - 1, connect[2]})
        edges[connect[1] - 1] = append(edges[connect[1] - 1], Edge{connect[0] - 1, connect[2]})
    }
    // 将与第一个城市连接的城市放入 heap 中
    for _, edge := range edges[0] {
        heap.Push(h,edge)
    }

    count := 1
    for h.Len() > 0 {    
        e := heap.Pop(h).(Edge)
        if visited[e.city] {
            continue
        }
        visited[e.city] = true
        for _, edge := range edges[e.city] {
            heap.Push(h, edge)
        }
        cost += e.cost
        count++
        if count == n {
            return cost
        }
    }
    return -1
}

type Edge struct {
    city int
    cost int
}

type EdgeHeap []Edge

func (h EdgeHeap) Len() int           { return len(h) }
func (h EdgeHeap) Less(i, j int) bool { return h[i].cost < h[j].cost }
func (h EdgeHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *EdgeHeap) Push(x interface{}) {
    *h = append(*h, x.(Edge))
}

func (h *EdgeHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}