题解 | #合并k个已排序的链表#

使用golang内置的"container/heap"实现优先队列后维护一个新的链表。

注意实现heap的五个接口：Len() int, Less(i, j int) bool, Swap(i, j int), Push(x interface{}), Pop() interface{}。

package main
import "fmt"
import "container/heap"
import . "nc_tools"
/*
 * type ListNode struct{
 *   Val int
 *   Next *ListNode
 * }
 */

type ListNodeHeap []*ListNode

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

func mergeKLists( lists []*ListNode ) *ListNode {
    if len(lists) == 0{
        fmt.Println()
        return nil
    }
    pq := &ListNodeHeap{}
    heap.Init(pq)
    for i := range lists{
        if lists[i] != nil{
            heap.Push(pq, lists[i])
        }
    }
    pHead := ListNode{}
    pCurr := &pHead
    for {
        if len(*pq) == 0{
            break
        }
        curr := heap.Pop(pq).(*ListNode)
        if curr.Next != nil{
            heap.Push(pq, curr.Next)
        }
        pCurr.Next = curr
        pCurr = pCurr.Next
        pCurr.Next = nil
    }
    return pHead.Next
}