题解 | #牛群的合并#

知识点：小顶堆

解题思路：用小顶推存储节点，由于小顶堆的特性，Pop出的就是最小的，将其加入结果res，同时判断其Next是否为空，不为空则放入小顶堆。一直重复上序步骤，知道小顶堆为空

package main

import "container/heap"
import . "nc_tools"

/*
 * type ListNode struct{
 *   Val int
 *   Next *ListNode
 * }
 */

/**
 * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
 *
 *
 * @param lists ListNode类一维数组
 * @return ListNode类
 */
type SmallHeap []*ListNode

func (h SmallHeap) Less(i, j int) bool  { return h[i].Val < h[j].Val }
func (h SmallHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h SmallHeap) Len() int            { return len(h) }
func (h *SmallHeap) Push(x interface{}) { *h = append(*h, x.(*ListNode)) }
func (h *SmallHeap) Pop() interface{} {
	n := len(*h)
	ans := (*h)[n-1]
	*h = (*h)[:n-1]
	return ans
}

func mergeKLists(lists []*ListNode) *ListNode {
	// write code here
	h := SmallHeap{}
	heap.Init(&h)
	for _, list := range lists {
		if list != nil {
			heap.Push(&h, list)
		}
	}
	res := &ListNode{}
	cur := res
	for len(h) > 0 {
		min := heap.Pop(&h).(*ListNode)
		cur.Next = min
		cur = cur.Next
		if min.Next != nil {
			heap.Push(&h, min.Next)
		}
	}
	return res.Next
}