23. Merge k Sorted Lists

merge-k-sorted-lists.go

// https://leetcode.com/problems/merge-k-sorted-lists
//
// Time: O(N log k)
// Space: O(k)
//
// k = len(lists)
// N = total number of nodes across all lists

// .................... //

/**
 * Definition for singly-linked list.
 * type ListNode struct {
 *     Val int
 *     Next *ListNode
 * }
 */

type minheap []*ListNode

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

func (h *minheap) Push(v any) { *h = append(*h, v.(*ListNode)) }

func (h *minheap) Pop() any {
    n := h.Len()
    res := (*h)[n-1]
    *h = (*h)[:n-1]
    return res
}

func mergeKLists(lists []*ListNode) *ListNode {
    head := &ListNode{}
    curr := head

    h := minheap{}
    for _, node := range lists {
       if node != nil {
           heap.Push(&h, node)
       }
    }

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

        curr.Next = &ListNode{Val: node.Val}
        curr = curr.Next

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

    return head.Next
}