import "container/heap"

func mergeKLists(lists []*ListNode) *ListNode {
    h := make(ListNodeHeap, 0)
    
    for _, n := range lists {
        if n != nil {
            h = append(h, n)
        }
    }
    
    heap.Init(&h)

    dummy := &ListNode{}
    tail := dummy

    for h.Len() > 0 {
        node := heap.Pop(&h).(*ListNode)
        
        tail.Next = node
        tail = node

        if node.Next != nil {
            heap.Push(&h, node.Next)
        }
    }

    return dummy.Next
}

// Implement heap using heap.Interface from 'container/heap' package.

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 any) { *h = append(*h, x.(*ListNode)) }
func (h *ListNodeHeap) Pop() any {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0:n-1]
    return x
}