https://leetcode.com/problems/top-k-frequent-words

top-k-frequent-words.go

import "container/heap"

type Item struct {
    word string
    count int
}

func (i Item) Less(o Item) bool {
    if i.count == o.count {
        return i.word > o.word
    }
    
    return i.count < o.count
}

type Heap []Item

func (h Heap) Len() int { return len(h) }

func (h Heap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

func (h Heap) Less(i, j int) bool { 
    return h[i].Less(h[j])
}

func (h *Heap) Push(x interface{}) {
    *h = append(*h, x.(Item))
}

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

func topKFrequent(words []string, k int) (result []string) {
    counts := make(map[string]int)
    
    for _, word := range words {
        counts[word]++
    }
    
    h := Heap{}
    
    for word, count := range counts {
        item := Item{word, count}
        
        if h.Len() < k {
            heap.Push(&h, item)
        } else if h[0].Less(item) {
            h[0] = item
            heap.Fix(&h, 0)
        }
    }
    
    for h.Len() > 0 {
        item := heap.Pop(&h).(Item)
        result = append(result, item.word)
    }
    
    reverse(result)
    
    return result
}

func reverse(arr []string) {
    for i, j := 0, len(arr) - 1; i < j; i, j = i + 1, j - 1 {
        arr[i], arr[j] = arr[j], arr[i]
    }
}