4.4.5 — Top K Frequent Elements

LeetCode 347 · Medium · ★ Blind 75

The problem

Return the k most frequent elements, in any order.

nums = [1,1,1,2,2,3], k = 2   →  [1, 2]

Up to 100,000 numbers. The answer is guaranteed unique, so there is no tie at the cutoff. The problem asks for something better than O(n \log n), and that requirement is the whole point.

The pattern

Two questions stacked. First, how often does each value appear — that is a counting map and it is the easy half. Second, which k counts are largest — that is a selection problem, and it has three answers worth knowing.

from collections import Counter
count = Counter(nums)          # {1: 3, 2: 2, 3: 1}

Solution 1 — sort the counts

count = Counter(nums)
return sorted(count.keys(), key=lambda x: count[x], reverse=True)[:k]

O(u \log u) for u distinct values. Correct, and the thing you are being asked to beat.

The waste: sorting ranks every value, and the question only asked for the top k.

Solution 2 — a heap of size k

Keep only k candidates and throw away anything that cannot make the cut.

import heapq

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        count = Counter(nums)
        heap = []
        for value, freq in count.items():
            heapq.heappush(heap, (freq, value))
            if len(heap) > k:
                heapq.heappop(heap)
        return [value for freq, value in heap]

Use a min-heap for the k largest. People get this backwards, so here is why. A heap gives cheap access to one end only, and heapq always puts the smallest item at the root. The item you most need to reach is the weakest of the ones you are keeping, because that is the one a new arrival will replace. A max-heap would put the strongest at the root, which is the item you never want to touch.

Pushing (freq, value) works because Python compares tuples element by element, so it orders by frequency first.

O(u \log k). Much better than sorting when k is small and u is large. It is also the only version that works on a stream, since it needs just k items in memory and never revisits an element.

Python has this built in: [v for v, _ in Counter(nums).most_common(k)]. Know it, and know it is a size-k heap underneath.

Solution 3 — bucket by frequency, in O(n)

Here is the observation that makes linear time possible:

A frequency can never be larger than n. With n elements, nothing appears more than n times.

So the counts are not arbitrary numbers — they are integers from 1 to n. And when the values you sort by are bounded small integers, you do not sort, you index. Make an array with one slot per possible frequency and drop each value into the slot matching its count.

Python

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        count = Counter(nums)

        buckets = [[] for _ in range(len(nums) + 1)]
        for value, freq in count.items():
            buckets[freq].append(value)

        result = []
        for freq in range(len(buckets) - 1, 0, -1):
            for value in buckets[freq]:
                result.append(value)
                if len(result) == k:
                    return result
        return result

TypeScript

function topKFrequent(nums: number[], k: number): number[] {
  const count = new Map<number, number>();
  for (const n of nums) count.set(n, (count.get(n) ?? 0) + 1);

  const buckets: number[][] = Array.from({ length: nums.length + 1 }, () => []);
  for (const [value, freq] of count) buckets[freq].push(value);

  const result: number[] = [];
  for (let freq = buckets.length - 1; freq >= 1; freq--) {
    for (const value of buckets[freq]) {
      result.push(value);
      if (result.length === k) return result;
    }
  }
  return result;
}

Why n+1 buckets. Index f holds every value that occurred exactly f times. If all n elements are identical, the count is n, so index n must exist. Slot 0 stays empty.

The line that will bite you. Do not write [[]] * (len(nums) + 1) in Python. That does not make n+1 lists — it makes one list and n+1 references to it, so appending to buckets[3] also appends to every other slot. The list comprehension evaluates [] fresh each time and gives you separate lists.

JavaScript has the same trap in a different costume: new Array(n).fill([]) fills every slot with the same array. Array.from({length: n}, () => []) calls the function per slot and is correct.

Why this is O(n). The outer loop runs n+1 times. The inner loop looks like it multiplies that, but it does not: each distinct value lives in exactly one bucket, so across the whole outer loop the inner loop visits each value once. Total inner work is O(u), not O(u) per iteration.

O(n) time, O(n) space.

Trace

nums = [1,1,1,2,2,3], k = 2. Counts are {1: 3, 2: 2, 3: 1}, so bucket 1 holds [3], bucket 2 holds [2], bucket 3 holds [1]. Walking down from 6: buckets 6, 5, 4 empty; bucket 3 gives 1; bucket 2 gives 2; the result has 2 items and returns [1, 2].

Which one to write

Present the ladder rather than jumping to the answer. Sorting is O(u \log u). A size-k min-heap is O(u \log k) and survives a stream. Bucketing is O(n) and meets the stated follow-up.

The honest caveat on bucketing: it allocates n+1 lists even when there are only three distinct values. With few distinct values and a huge n, the heap uses far less memory.

Edge cases

If k equals the number of distinct values, the bucket walk never hits the early return and falls out at the end — which is why that final return result is not dead code. Negative values are fine; they are dictionary keys, and only the frequencies ever index the bucket array.

Where this goes next

• Top K Frequent Words — same problem with a lexicographic tie-break, which breaks the neatness of both fast solutions because order within a bucket now matters.
• Kth Largest Element in an Array — selection without the counting step, and it has a fourth solution: quickselect, O(n) on average. Chapter 4.17.
• Heavy hitters in a stream too large to count exactly — a Count-Min sketch. 4.29.

The rule: when the key you order by is a bounded small integer, do not sort — index. Counting sort, radix sort and this problem are one idea.

What the interviewer will push on

"Can you beat O(n \log n)?" They want the sentence "a frequency cannot exceed n, so counts are bounded small integers and can be array indices."

"Why a min-heap for the k largest?" The root must be the item you are about to evict.

"What if k is close to u?" The heap degenerates to O(u \log u) and stops helping. Bucketing does not care.

"What if the data is a stream?" Only the heap survives. If even the counting map will not fit, Count-Min sketch.

One thing to volunteer: name all three solutions and the condition that picks each one, in fifteen seconds, before writing anything.

Next: 4.4.6 Encode and Decode Strings — a design question rather than a counting one.