Top K Frequent Elements (using Heap) - LeetCode 347

top_k_frequent_w_heap.gol

type MaxHeap [][]int

func (h MaxHeap) Len() int { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i][0] > h[j][0] }
func (h MaxHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

func (h *MaxHeap) Push(x interface{}) {
    *h = append(*h, x.([]int))
}

func (h *MaxHeap) Pop() interface{} {
    old := *h
    n := len(old)
    val := old[n-1]
    *h = old[:n-1]
    return val
}

func topKFrequent(nums []int, k int) []int {
    freqMap := make(map[int]int)
    for _, n := range nums {
        freqMap[n]++
    }
    
    
    h := &MaxHeap{}
    var result []int
    
    for k, v := range freqMap {
        heap.Push(h, []int{v, k})
    }
    
    for i := 0; i < k; i++ {
        val := heap.Pop(h).([]int)
        result = append(result, val[1])
    }
    
    return result
}