LC23 - 合并K个升序链表

package LinkedList

import (
    "container/heap"
)

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{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

// 时间复杂度: O(n * log k), 空间复杂度: O(k)
// n表示节点总数, k表示链表个数
func mergeKLists(lists []*ListNode) *ListNode {
    minHeap := &ListNodeHeap{}
    heap.Init(minHeap)
    for _, list := range lists {
        if list != nil {
            heap.Push(minHeap, list)
        }
    }
    virtualHead := new(ListNode)
    tail := virtualHead
    for minHeap.Len() > 0 {
        cur := heap.Pop(minHeap).(*ListNode)
        tail.Next = cur
        tail = tail.Next
        if cur.Next != nil {
            heap.Push(minHeap, cur.Next)
        }
    }
    return virtualHead.Next
}

本题与LC21没有太大区别，唯一需要注意的是堆结构的实现。