https://leetcode.com/problems/smallest-range-covering-elements-from-k-lists/

smallest-range-covering-elements-from-k-lists.go

import "container/heap"

type HeapItem struct {
    value int
    listID int
}

type Heap []HeapItem

func (h Heap) Len() int { return len(h) }
func (h Heap) Less(i, j int) bool { return h[i].value < h[j].value }
func (h Heap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *Heap) Push(x interface{}) { *h = append(*h, x.(HeapItem)) }

func (h *Heap) Pop() interface{} { 
    n := h.Len()
    old := *h
    x := old[n - 1]
    *h = old[:n - 1]
    
    return x
}

func smallestRange(nums [][]int) []int {
    k := len(nums)
    r := []int{(1 << 31) - 1, -1 << 31}
    h := Heap{}
    
    for i, list := range nums {
        heap.Push(&h, HeapItem{
            value: list[0],
            listID: i,
        })
        
        r[0] = min(r[0], list[0])
        r[1] = max(r[1], list[0])
    }
    
    result := []int{r[0], r[1]}
    
    for h.Len() == k {
        item := h[0]
        nums[item.listID] = nums[item.listID][1:]
        
        if (r[1] - r[0]) < (result[1] - result[0]) {
            result[0], result[1] = r[0], r[1]
        }
        
        if len(nums[item.listID]) == 0 {
            break
        }
        
        h[0].value = nums[item.listID][0]
        heap.Fix(&h, 0)
        
        r[0] = h[0].value
        r[1] = max(r[1], nums[item.listID][0])
    }
    
    return result
}

func min(a, b int) int {
    if a < b {
        return a
    }
    
    return b
}

func max(a, b int) int {
    if a < b {
        return b
    }
    
    return a
}