cabloo · June 17, 2021 18:44 · Jun 17, 2021
diff --git a/heap.go b/heap.go
@@ -0,0 +1,68 @@
+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
+}
No results found