刷题使我快乐,满脸开心.jpg

  • 来源:力扣(LeetCode)
  • 链接:https://leetcode.cn/problems/top-k-frequent-words/
  • 著作权归领扣网络所有。商业转载请联系官方授权,非商业转载请注明出处。

题目

给定一个单词列表 words 和一个整数 k ,返回前 k 个出现次数最多的单词。
返回的答案应该按单词出现频率由高到低排序。如果不同的单词有相同出现频率, 按字典顺序 排序。

示例 1:

输入: words = ["i", "love", "leetcode", "i", "love", "coding"], k = 2
输出: ["i", "love"]
解析: "i" 和 "love" 为出现次数最多的两个单词,均为2次。
    注意,按字母顺序 "i" 在 "love" 之前。

示例 2:

输入: ["the", "day", "is", "sunny", "the", "the", "the", "sunny", "is", "is"], k = 4
输出: ["the", "is", "sunny", "day"]
解析: "the", "is", "sunny" 和 "day" 是出现次数最多的四个单词,
    出现次数依次为 4, 3, 2 和 1 次。

注意:

  • 1 <= words.length <= 500
  • 1 <= words[i] <= 10
  • words[i] 由小写英文字母组成。
  • k 的取值范围是 [1, 不同 words[i] 的数量]

思路

top k 的小变种,先计算下权值,然后搞一个优先队列即可,不过这里有个优化点

因为要保留的只是最大的k个值,所以可以使用小顶堆时,可以在元素个数超过k后直接pop掉多余的值,减少堆调整的消耗。只是需要记得输出结果时,pop出来的顺序是反的

代码

func topKFrequent(words []string, k int) []string {
	workCountMap := make(map[string]int)
	for _, work := range words {
		workCountMap[work]++
	}
	intHeap := &IntHeap{}
	for word, count := range workCountMap {
		heap.Push(intHeap, pair{
			Word:  word,
			Count: count,
		})
		if intHeap.Len() > k {
			heap.Pop(intHeap)
		}
	}
	res := make([]string, k)
	for i := k - 1; i >= 0; i-- {
		res[i] = heap.Pop(intHeap).(pair).Word
	}
	return res
}

type IntHeap []pair

type pair struct {
	Word  string
	Count int
}

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

func (h IntHeap) Less(i, j int) bool {
	return h[i].Count < h[j].Count || (h[i].Count == h[j].Count && h[i].Word > h[j].Word)
}

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

func (h *IntHeap) Pop() interface{} {
	old := *h
	v := old[len(old)-1]
	*h = old[:len(old)-1]
	return v
}

func (h *IntHeap) Push(v interface{}) {
	*h = append(*h, v.(pair))
}