heap.go

import "container/heap"

type CountPriorityQueueItem struct {
    Value, Count, index int
}

type CountPriorityQueue struct {
    items []*CountPriorityQueueItem
    byValue map[int]*CountPriorityQueueItem
}

func (q *CountPriorityQueue) Pop() interface{} {
    last := len(q.items)-1
    item := q.items[last]
    q.items[last] = nil
    q.items = q.items[:last]
    return *item
}

func (q *CountPriorityQueue) PopItem() CountPriorityQueueItem {
    return heap.Pop(q).(CountPriorityQueueItem)
}

func (q *CountPriorityQueue) Push(x interface{}) {
    value := x.(int)
    if curr, ok := q.byValue[value]; ok {
        curr.Count++
        heap.Fix(q, curr.index)
        return
    }
    item := &CountPriorityQueueItem{Value: value, Count: 1, index: len(q.items)}
    q.byValue[value] = item
    q.items = append(q.items, item)
    heap.Fix(q, item.index)
}

func (q CountPriorityQueue) Len() int {
    return len(q.items)
}

func (q CountPriorityQueue) Swap(i, j int) {
    q.items[i], q.items[j] = q.items[j], q.items[i]
    q.items[i].index = i
    q.items[j].index = j
}

func (q CountPriorityQueue) Less(a, b int) bool {
    return q.items[a].Count >= q.items[b].Count
}

func topKFrequent(nums []int, k int) []int {
    maxCountHeap := &CountPriorityQueue{nil, map[int]*CountPriorityQueueItem{}}
    for _, v := range nums {
        maxCountHeap.Push(v)
    }
    heap.Init(maxCountHeap)
    
    for _, v := range maxCountHeap.items {
        fmt.Println(v.Value, v.Count, v.index)
    }
    
    result := []int{}
    for ; k > 0; k-- {
        result = append(result, maxCountHeap.PopItem().Value)
    }
    
    return result
}