Top K Frequent Words

Medium Heaps Top K / Kth Largest / Smallest Original on LeetCode

Reported atAmazonGoldman Sachs

Top K Frequent Words is a medium heaps problem solved with the top k / kth largest / smallest pattern. The best approach, optimal (min-heap of size k), runs in O(n + m log k) time and O(m) space. Below are 2 approaches in Java, Python, C++ and JavaScript, from brute force (count, then sort all words) up.

Problem

You are given an array of strings words and an integer k. Return the k words that occur most often in words.

Order the result from the highest count to the lowest. When two words have the same count, the one that comes first in alphabetical (dictionary) order goes first.

Examples

Example 1

Input
words = ["tea", "jam", "tea", "bun", "jam", "tea"], k = 2
Output
["tea", "jam"]
Why
"tea" appears 3 times and "jam" twice, so they are the two most frequent words.

Example 2

Input
words = ["pear", "fig", "kiwi", "fig", "pear", "kiwi", "date"], k = 3
Output
["fig", "kiwi", "pear"]
Why
Three words tie with 2 occurrences each, so they are listed in alphabetical order.

Example 3

Input
words = ["sun", "moon", "star"], k = 1
Output
["moon"]
Why
Every word appears once; "moon" comes first alphabetically.

Constraints

  • 1 <= words.length <= 10^4
  • 1 <= words[i].length <= 10; words contain only lowercase English letters.
  • 1 <= k <= number of distinct words.
  • Follow-up: can you do it in O(n log k) time?

Solve it here

Write your solution in Java, Python, C++ or JavaScript and run it against the examples, right in the page. Submit runs the hidden tests too and, with a free account, records the solve in your revision log.

00:00
Testcase

Run checks the testcases above (3 examples, or your own). Submit also runs 11 hidden tests and saves the solve to your log. Signed in, every Run and Submit is saved under My submissions; your draft stays in this browser.

Solutions

Try it yourself first. Then compare: each approach lists its idea, the steps, its time and space, and the code in Java, Python, C++ and JavaScript.

ApproachTimeSpace
Brute force (count, then sort all words)O(n + m log m)O(m)
Optimal (min-heap of size k)O(n + m log k)O(m)

1Brute force (count, then sort all words)

TimeO(n + m log m)Counting is O(n); sorting the m distinct words costs O(m log m) comparisons (each comparison is at most 10 characters).
SpaceO(m)The map and the list hold m distinct words.

Count every word, put the distinct words in a list, and sort that list by count descending, breaking ties alphabetically. The first k entries are the answer.

  1. Build a map word -> count.
  2. Copy the distinct words into a list.
  3. Sort: if two counts differ, the larger count goes first; otherwise compare the words alphabetically.
  4. Return the first k words.
class Solution {
    public List<String> topKFrequent(String[] words, int k) {
        Map<String, Integer> count = new HashMap<>();
        for (String w : words) count.merge(w, 1, Integer::sum);
        List<String> distinct = new ArrayList<>(count.keySet());
        distinct.sort((a, b) -> {
            int ca = count.get(a), cb = count.get(b);
            return ca != cb ? cb - ca : a.compareTo(b);
        });
        return new ArrayList<>(distinct.subList(0, k));
    }
}

2Optimal (min-heap of size k)

TimeO(n + m log k)Counting is O(n); each of the m distinct words costs one push and at most one pop on a heap of size k.
SpaceO(m)The count map holds m words; the heap holds at most k + 1.

After counting, scan the distinct words and keep only the best k seen so far in a min-heap. The heap's top is the weakest of them: the smallest count, and among equal counts the alphabetically larger word. When the heap grows past k, drop its top. At the end, pop everything (weakest first) and reverse to get the best first.

  1. Build a map word -> count.
  2. Create a heap whose top is the weakest word: lower count is weaker; on equal counts, the alphabetically larger word is weaker.
  3. For each distinct word: push it; if the heap size exceeds k, pop the top.
  4. Pop all k words into a list (weakest first) and reverse it.
class Solution {
    public 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) -> {
            int ca = count.get(a), cb = count.get(b);
            return ca != cb ? ca - cb : b.compareTo(a);
        });
        for (String w : count.keySet()) {
            heap.offer(w);
            if (heap.size() > k) heap.poll();
        }
        List<String> result = new ArrayList<>();
        while (!heap.isEmpty()) result.add(heap.poll());
        Collections.reverse(result);
        return result;
    }
}

Edge cases to test

  • All words are different, so the answer is the first k words in alphabetical order
  • Ties in count that straddle the cut-off at position k
  • k equals the number of distinct words
  • Words where one is a prefix of another, like "ab" and "abc": the shorter one comes first

Hints

Hint 1

Start by counting how often each word appears with a hash map.

Hint 2

The required order is: higher count first, and for equal counts, alphabetical. Sorting with that comparator already solves it.

Hint 3

To do better than sorting everything, keep a min-heap of size k whose top is the weakest candidate: the lowest count, and on a tie the alphabetically larger word.

FAQ

What is the best time complexity for Top K Frequent Words?

Optimal (min-heap of size k) runs in O(n + m log k) time and O(m) extra space. Counting is O(n); each of the m distinct words costs one push and at most one pop on a heap of size k.

Which pattern does Top K Frequent Words use?

It is a heaps problem that uses the top k / kth largest / smallest pattern. Other problems with the same pattern: Kth Largest Element in an Array, Kth Largest Element in a Stream.

Is there a brute force solution for Top K Frequent Words?

Yes. Brute force (count, then sort all words) takes O(n + m log m) time and O(m) space. Count every word, put the distinct words in a list, and sort that list by count descending, breaking ties alphabetically.

Which edge cases should I test for Top K Frequent Words?

All words are different, so the answer is the first k words in alphabetical order; Ties in count that straddle the cut-off at position k; k equals the number of distinct words; Words where one is a prefix of another, like "ab" and "abc": the shorter one comes first.