leetcode-go/leetcode/medium/692.top-k-frequent-words.go

type pair struct {
	word string
	freq int
}

type minHeap []pair

func (h minHeap) Len() int { return len(h) }
func (h minHeap) Less(i, j int) bool {
	if h[i].freq != h[j].freq {
		return h[i].freq < h[j].freq
	}
	return h[j].word < h[i].word
}
func (h minHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

func (h *minHeap) Push(x interface{}) {
	*h = append(*h, x.(pair))
}

func (h *minHeap) Pop() interface{} {
	l := h.Len()
	x := (*h)[l-1]
	*h = (*h)[:l-1]
	return x
}

func topKFrequent(words []string, k int) []string {
	m := make(map[string]int)
	for _, word := range words {
		m[word]++
	}
	i := 0
	var h minHeap
	for word, freq := range m {
		heap.Push(&h, pair{word, freq})
		i++
		if k < i {
			heap.Pop(&h)
		}
	}
	res := make([]string, k)
	for i = k - 1; 0 <= i; i-- {
		p := heap.Pop(&h).(pair)
		res[i] = p.word
	}
	return res
}