692. 前 K 个高频单词

题目描述

输入单词数组 words 和整数 k。

统计每个不同单词出现的次数，选出次数最多的 k 个单词返回。

结果按出现次数从多到少排列；次数相同的单词按字典序从小到大排列。

示例 1

输入：words = ["i","love","leetcode","i","love","coding"]
k = 2
输出：["i","love"]

解释："i" 和 "love" 各出现 2 次，按字典序 "i" 在前。

示例 2

输入：words = ["the","day","is","sunny","the","the","the","sunny","is","is"]
k = 4
输出：["the","is","sunny","day"]

解释：四个单词的出现次数依次为 4、3、2、1。

数据范围

1 <= words.length <= 500
1 <= words[i].length <= 10
words[i] 只包含小写英文字母。
k 的取值范围是 [1, 不同单词的个数]。

进阶：能否做到时间 O(n log k)、空间 O(n)？

解题思路：哈希表 + 堆

先用哈希表统计每个单词的次数。然后把 (−次数, 单词) 放进最小堆：次数越大越靠前，次数相同时元组自然比较单词，字典序小的靠前，恰好就是题目要求的顺序。heapify 建堆是 O(m)，再弹出 k 次就是答案。

为什么正确

堆顶始终是剩余单词中“次数最大，其次字典序最小”的那一个，所以依次弹出的就是排好序的结果。

复杂度分析

时间 O(n + m + k log m)，n 为单词总数，m 为不同单词数；空间 O(m)。

易错点

次数取负是为了用最小堆模拟“次数大的优先”；但单词不能取负，字典序小的本来就应该优先。若用大小为 k 的堆优化，比较规则要完全反过来（次数小、字典序大的先被淘汰），最后还要反转。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class Solution:
    def topKFrequent(self, words: List[str], k: int) -> List[str]:
        cnt = collections.Counter(words)
        heap = [(-c, w) for w, c in cnt.items()]
        heapq.heapify(heap)
        res = []
        for _ in range(k):
            res.append(heapq.heappop(heap)[1])
        return res

class Solution {
    public List<String> topKFrequent(String[] words, int k) {
        Map<String, Integer> cnt = new HashMap<>();
        for (String w : words) cnt.merge(w, 1, Integer::sum);
        PriorityQueue<String> heap = new PriorityQueue<>((x, y) ->
            !cnt.get(x).equals(cnt.get(y)) ? cnt.get(y) - cnt.get(x) : x.compareTo(y));
        heap.addAll(cnt.keySet());
        List<String> res = new ArrayList<>();
        for (int i = 0; i < k; i++) {
            res.add(heap.poll());
        }
        return res;
    }
}

class Solution {
public:
    vector<string> topKFrequent(vector<string>& words, int k) {
        unordered_map<string, int> cnt;
        for (auto& w : words) cnt[w]++;
        auto worse = [](const pair<int, string>& x, const pair<int, string>& y) {
            return x.first != y.first ? x.first < y.first : x.second > y.second;
        };
        priority_queue<pair<int, string>, vector<pair<int, string>>, decltype(worse)> heap(worse);
        for (auto& [w, c] : cnt) heap.push({c, w});
        vector<string> res;
        for (int i = 0; i < k; i++) {
            res.push_back(heap.top().second);
            heap.pop();
        }
        return res;
    }
};

type wc struct {
	w string
	c int
}
type wordHeap []wc

func (h wordHeap) Len() int { return len(h) }
func (h wordHeap) Less(i, j int) bool {
	if h[i].c != h[j].c {
		return h[i].c > h[j].c
	}
	return h[i].w < h[j].w
}
func (h wordHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *wordHeap) Push(x any)   { *h = append(*h, x.(wc)) }
func (h *wordHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}

func topKFrequent(words []string, k int) []string {
	cnt := map[string]int{}
	for _, w := range words {
		cnt[w]++
	}
	h := &wordHeap{}
	for w, c := range cnt {
		*h = append(*h, wc{w, c})
	}
	heap.Init(h)
	res := []string{}
	for i := 0; i < k; i++ {
		res = append(res, heap.Pop(h).(wc).w)
	}
	return res
}

class MinHeap {
    constructor(less) { this.a = []; this.less = less; }
    get size() { return this.a.length; }
    push(x) {
        const a = this.a; a.push(x);
        for (let i = a.length - 1; i > 0;) {
            const p = (i - 1) >> 1;
            if (!this.less(a[i], a[p])) break;
            [a[i], a[p]] = [a[p], a[i]]; i = p;
        }
    }
    pop() {
        const a = this.a, top = a[0], last = a.pop();
        if (a.length) {
            a[0] = last;
            for (let i = 0; ;) {
                let m = i; const l = 2 * i + 1, r = l + 1;
                if (l < a.length && this.less(a[l], a[m])) m = l;
                if (r < a.length && this.less(a[r], a[m])) m = r;
                if (m === i) break;
                [a[i], a[m]] = [a[m], a[i]]; i = m;
            }
        }
        return top;
    }
}

var topKFrequent = function(words, k) {
    const cnt = new Map();
    for (const w of words) cnt.set(w, (cnt.get(w) || 0) + 1);
    const heap = new MinHeap((x, y) => x[0] !== y[0] ? x[0] > y[0] : x[1] < y[1]);
    for (const [w, c] of cnt) heap.push([c, w]);
    const res = [];
    for (let i = 0; i < k; i++) {
        res.push(heap.pop()[1]);
    }
    return res;
};

class MinHeap<T> {
    private a: T[] = [];
    constructor(private less: (x: T, y: T) => boolean) {}
    get size(): number { return this.a.length; }
    push(x: T): void {
        const a = this.a; a.push(x);
        for (let i = a.length - 1; i > 0;) {
            const p = (i - 1) >> 1;
            if (!this.less(a[i], a[p])) break;
            [a[i], a[p]] = [a[p], a[i]]; i = p;
        }
    }
    pop(): T {
        const a = this.a, top = a[0], last = a.pop()!;
        if (a.length) {
            a[0] = last;
            for (let i = 0; ;) {
                let m = i; const l = 2 * i + 1, r = l + 1;
                if (l < a.length && this.less(a[l], a[m])) m = l;
                if (r < a.length && this.less(a[r], a[m])) m = r;
                if (m === i) break;
                [a[i], a[m]] = [a[m], a[i]]; i = m;
            }
        }
        return top;
    }
}

function topKFrequent(words: string[], k: number): string[] {
    const cnt = new Map<string, number>();
    for (const w of words) cnt.set(w, (cnt.get(w) || 0) + 1);
    const heap = new MinHeap<[number, string]>((x, y) => x[0] !== y[0] ? x[0] > y[0] : x[1] < y[1]);
    for (const [w, c] of cnt) heap.push([c, w]);
    const res: string[] = [];
    for (let i = 0; i < k; i++) {
        res.push(heap.pop()![1]);
    }
    return res;
}

载入编辑器，自己跑一遍 →

其他解法

两种解法都先用哈希表计数，比较规则相同；排序一次排好全部单词，堆只弹出前 k 个。

解法 2：哈希表 + 排序时间 O(n log n)，空间 O(n)

先用哈希表统计每个单词的出现次数。再把所有不同单词排序：次数多的在前，次数相同时字典序小的在前。取前 k 个。

只需写好一个比较规则，代码最短。代价是对全部不同单词排序，k 远小于单词数时做了多余的工作。

class Solution:
    def topKFrequent(self, words: List[str], k: int) -> List[str]:
        cnt = collections.Counter(words)
        return sorted(cnt, key=lambda w: (-cnt[w], w))[:k]

class Solution {
    public List<String> topKFrequent(String[] words, int k) {
        Map<String, Integer> cnt = new HashMap<>();
        for (String w : words) cnt.merge(w, 1, Integer::sum);
        List<String> keys = new ArrayList<>(cnt.keySet());
        keys.sort((a, b) -> {
            int ca = cnt.get(a), cb = cnt.get(b);
            return ca != cb ? cb - ca : a.compareTo(b);
        });
        return keys.subList(0, k);
    }
}

class Solution {
public:
    vector<string> topKFrequent(vector<string>& words, int k) {
        unordered_map<string, int> cnt;
        for (auto& w : words) cnt[w]++;
        vector<string> keys;
        for (auto& [w, c] : cnt) keys.push_back(w);
        sort(keys.begin(), keys.end(), [&](const string& a, const string& b) {
            return cnt[a] != cnt[b] ? cnt[a] > cnt[b] : a < b;
        });
        keys.resize(k);
        return keys;
    }
};

func topKFrequent(words []string, k int) []string {
	cnt := map[string]int{}
	for _, w := range words {
		cnt[w]++
	}
	keys := make([]string, 0, len(cnt))
	for w := range cnt {
		keys = append(keys, w)
	}
	sort.Slice(keys, func(i, j int) bool {
		a, b := keys[i], keys[j]
		if cnt[a] != cnt[b] {
			return cnt[a] > cnt[b]
		}
		return a < b
	})
	return keys[:k]
}

var topKFrequent = function(words, k) {
    const cnt = new Map();
    for (const w of words) cnt.set(w, (cnt.get(w) || 0) + 1);
    const keys = [...cnt.keys()];
    keys.sort((a, b) => {
        const ca = cnt.get(a), cb = cnt.get(b);
        if (ca !== cb) return cb - ca;
        return a < b ? -1 : a > b ? 1 : 0;
    });
    return keys.slice(0, k);
};

function topKFrequent(words: string[], k: number): string[] {
    const cnt = new Map<string, number>();
    for (const w of words) cnt.set(w, (cnt.get(w) || 0) + 1);
    const keys = [...cnt.keys()];
    keys.sort((a, b) => {
        const ca = cnt.get(a)!, cb = cnt.get(b)!;
        if (ca !== cb) return cb - ca;
        return a < b ? -1 : a > b ? 1 : 0;
    });
    return keys.slice(0, k);
}

栈与堆专题的其他题目

查看栈与堆全部题目 →