Top K Frequent Words

mediumHeap / Priority Queue38 stepsJava · Python · C++

Return the k most frequent words, ties broken alphabetically. Count words, then keep a size-k min-heap whose "smallest" is the least frequent word (and, among equal counts, the alphabetically last). Pop everything and reverse.

▶ Watch the animated trace

Traced input: words = [i, love, leetcode, i, love, coding], k = 2 → ["i", "love"]. Original problem: leetcode.com/problems/top-k-frequent-words/

Complexity

BestO(n log k)

AverageO(n log k)

WorstO(n log k)

SpaceO(n)

n distinct words.

Key idea

The heap's top must be the word that should be evicted first.
Evict lower counts first; for equal counts, evict the alphabetically larger word.
Popping gives worst-to-best, so reverse at the end.

Common mistakes

Getting the tie-break backwards in the min-heap comparator.
Forgetting to reverse the popped order.

Solution

Java

List<String> topKFrequent(String[] words, int k) {
    Map<String, Integer> count = new HashMap<>();
    for (String w : words) count.merge(w, 1, Integer::sum);
    PriorityQueue<String> heap = new PriorityQueue<>((a, b) -> count.get(a).equals(count.get(b)) ? b.compareTo(a) : count.get(a) - count.get(b));
    for (String w : count.keySet()) {
        heap.offer(w);
        if (heap.size() > k) {
            heap.poll();
        }
    }
    LinkedList<String> res = new LinkedList<>();
    while (!heap.isEmpty()) res.addFirst(heap.poll());
    return res;
}

Python

class Word:
    def __init__(self, word: str, count: int):
        self.word, self.count = word, count

    def __lt__(self, other: "Word") -> bool:
        # the heap's top is the word to evict first: lower count, then alphabetically later
        if self.count != other.count: return self.count < other.count
        return self.word > other.word


def topKFrequent(words: List[str], k: int) -> List[str]:
    count = Counter(words)
    heap = []
    for w, c in count.items():
        heapq.heappush(heap, Word(w, c))
        if len(heap) > k:
            heapq.heappop(heap)
    res = []
    while heap: res.append(heapq.heappop(heap).word)
    return res[::-1]

C++

vector<string> topKFrequent(vector<string>& words, int k) {
    unordered_map<string, int> count;
    for (auto& w : words) count[w]++;
    // top of the heap = the word to evict first: lower count, then alphabetically later
    auto cmp = [&](const string& a, const string& b) {
        return count[a] != count[b] ? count[a] > count[b] : a < b;
    };
    priority_queue<string, vector<string>, decltype(cmp)> heap(cmp);
    for (auto& [w, c] : count) {
        heap.push(w);
        if ((int)heap.size() > k) {
            heap.pop();
        }
    }
    vector<string> res;
    while (!heap.empty()) { res.push_back(heap.top()); heap.pop(); }
    reverse(res.begin(), res.end());
    return res;
}

Edge cases

Each one has its own trace in the interactive view.

Case	Input	Answer
Ties broken alphabetically All counts are 1, so alphabetical order decides.	`words = [b, a, c], k = 2`	`["a", "b"]`
k = 1 y appears twice.	`words = [x, y, y], k = 1`	`["y"]`
LeetCode example 2 Counts 4, 3, 2, 1.	`words = [the, day, is, sunny, the, the, the, sunny, is, is], k = 4`	`["the", "is", "sunny", "day"]`