LeetCode 692: Top K Frequent Words

Sorted frequency map

We first create the frequency map: {word: count}. We then sort the map items with (- count) as the primary key and word as the secondary key. Finally, we return the first k of these sorted words as the answer.

Assuming word comparison cost as constant, If there are $n$ words, time: $\mathcal{O}(n \cdot \lg{n})$ , space: $\mathcal{O}(n)$ .

class Solution:
    def topKFrequent(self, words: List[str], k: int) -> List[str]:
        freq_map = {}
        for w in words:
            freq_map[w] = freq_map.get(w, 0) + 1
        sorted_words = [w for w, _ in sorted(freq_map.items(), key=lambda x: (-x[1], x[0]))]
        return sorted_words[:k]

Follow-up: Could you solve it in O(n log(k)) time and O(n) extra space?

We can use a min heap with custom comparison so that a tie in frequency count is broken in favor of bigger words.

Assuming $n > k$ , time: $\mathcal{O}(n \cdot \lg{k} + k \cdot \lg{k}) = \mathcal{O}(n \cdot \lg{k})$ , space: $\mathcal{O}(n + k) = \mathcal{O}(n)$ .

from heapq import heappush as push, heappop as pop


class HeapItem:
    def __init__(self, count, word):
        self.count = count
        self.word = word

    def __lt__(self, other):
        if self.count != other.count:
            return self.count < other.count
        return self.word >= other.word

class Solution:
    def topKFrequent(self, words: List[str], k: int) -> List[str]:
        freq_map = {}
        for w in words:
            freq_map[w] = freq_map.get(w, 0) + 1

        q = []
        for word, count in freq_map.items():
            push(q, HeapItem(count, word))
            if len(q) > k:
                pop(q)

        q.sort(reverse=True)
        return [item.word for item in q]

unreasonably effective