347. Top K Frequent Elements

top-k-frequent-elements.go

// https://leetcode.com/problems/top-k-frequent-elements/
//
// Time: O(n log k)
// Space: O(n)
//
// n - number of unique elements in array

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

type item struct {
    num int
    freq int
}

type itemheap []item

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

func (h *itemheap) Push(v any) {
    *h = append(*h, v.(item))
}

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

func topKFrequent(nums []int, k int) []int {
    freqByNum := make(map[int]int)
    for _, num := range nums {
        freqByNum[num]++
    }

    h := itemheap{}
    for num, freq := range freqByNum {
        heap.Push(&h, item{num, freq})
        if h.Len() > k {
            _ = heap.Pop(&h)
        }
    }

    res := make([]int, k)
    for i := range k {
        res[k-i-1] = heap.Pop(&h).(item).num
    }

    return res
}