692. Top K Frequent Words

top-k-frequent-words.go

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

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

type item struct {
    word string
    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 {
    if h[i].freq == h[j].freq {
        return h[i].word > h[j].word
    }
    return h[i].freq < h[j].freq 
}

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(words []string, k int) []string {
    freqByWord := make(map[string]int)
    for _, word := range words {
        freqByWord[word]++
    }

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

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

    return res
}