LeetCode 347 – Top K Frequent Elements

The problem statement contains its own hint: "your algorithm's time complexity must be better than O(n log n)". That sentence exists to rule out the two solutions everybody reaches for first — sort by frequency, or push everything into a heap — and it is the reason this question is asked at all.

The problem

Given an integer array and a number k, return the k most frequent elements, in any order. k is always valid.

nums = [1,1,1,2,2,3], k = 2  -> [1, 2]
nums = [1],           k = 1  -> [1]
nums = [1,2],         k = 2  -> [1, 2]   ties, both returned

The heap trap

Counting frequencies is obviously step one. What you do with the counts is the question, and there is a specific wrong turn worth naming because it is so easy to defend badly:

A max-heap of everything is O(n log n). The reasoning that makes it feel fine goes: "there are only m unique values, and m is usually much smaller than n." Usually. But [1, 2, 3, …, n] has m = n, so building the heap is n insertions at O(log n) each — exactly the bound the problem forbids. An argument that relies on the input being friendly is not a complexity argument.

A min-heap of size k is O(m log k), and that is fine. Keep only k entries; when a new one arrives and the heap is full, compare against the smallest and evict if the newcomer beats it. Since k ≤ m, log k is never worse than log m, and this satisfies the constraint. It is a perfectly good answer.

But there is a better one, because the counts have a property that lets you skip comparison sorting entirely.

Bucket sort: the O(n) answer

A frequency is bounded. No value can occur more than n times, and every count is a positive integer. So the counts can be used directly as array indices — no comparisons needed.

Build an array of n + 1 buckets where buckets[f] holds every value that occurred exactly f times. Then walk it from the high end down, collecting until you have k:

nums = [1,1,1,2,2,3]      counts: 1 -> 3,  2 -> 2,  3 -> 1

bucket:  0    1    2    3    4    5    6
        [ ]  [3]  [2]  [1]  [ ]  [ ]  [ ]
                        ↑ walk down from 6, take 1, then 2, stop at k = 2

Most buckets are empty and that is fine — the array is O(n) to allocate and O(n) to walk, and skipping empties costs nothing. This is counting sort's core idea: when the keys are small bounded integers, indexing beats comparing.

Java

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> freq = new HashMap<>();
        for (int num : nums) {
            freq.merge(num, 1, Integer::sum);
        }

        // buckets[f] holds every value seen exactly f times.
        // Nothing can occur more than nums.length times, so this is wide enough.
        List<Integer>[] buckets = new List[nums.length + 1];
        for (Map.Entry<Integer, Integer> entry : freq.entrySet()) {
            int count = entry.getValue();
            if (buckets[count] == null) {
                buckets[count] = new ArrayList<>();
            }
            buckets[count].add(entry.getKey());
        }

        int[] result = new int[k];
        int filled = 0;
        for (int f = nums.length; f >= 1 && filled < k; f--) {
            if (buckets[f] == null) continue;

            for (int value : buckets[f]) {
                if (filled == k) break;      // a bucket can overshoot k on a tie
                result[filled++] = value;
            }
        }

        return result;
    }
}

freq.merge(num, 1, Integer::sum) is the tidy way to write "insert 1 or add 1"; getOrDefault(num, 0) + 1 reads more plainly if you prefer. Either beats containsKey followed by get, which hashes twice.

The inner filled == k break is not redundant with the outer condition. A single bucket can hold several values, and if k = 1 with two values tied at the top frequency, the outer loop only re-checks between buckets — the inner break is what stops the array overrunning inside one.

new List[…] raises an unchecked warning; generic array creation is not allowed in Java. It is the accepted idiom here, and saying so pre-empts the question.

Python

from collections import Counter


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

        # buckets[f] holds every value seen exactly f times.
        buckets = [[] for _ in range(len(nums) + 1)]
        for value, count in freq.items():
            buckets[count].append(value)

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

        return result

Counter(nums).most_common(k) solves the whole thing in one line, and internally uses the size-k heap — so it is O(n log k), not O(n). Mention it, then write the buckets.

The heap version, for completeness

// O(m log k). Correct, and the one to write if asked for O(1)-ish extra space.
PriorityQueue<Map.Entry<Integer, Integer>> heap =
    new PriorityQueue<>(Map.Entry.comparingByValue());   // MIN-heap on frequency

for (Map.Entry<Integer, Integer> entry : freq.entrySet()) {
    heap.offer(entry);
    if (heap.size() > k) heap.poll();    // evict the least frequent
}

Min-heap, not max-heap. The counter-intuitive part is that you evict the smallest to keep the largest, and the heap's root is therefore the weakest survivor rather than the answer.

Complexity

Approach	Time	Space
Sort by frequency	`O(n log n)` — violates the constraint	`O(n)`
Max-heap of everything	`O(n log n)` — violates it too	`O(n)`
Min-heap of size k	`O(m log k)`	`O(m + k)`
Bucket sort	`O(n)`	`O(n)`

Bucket sort trades space for time — the array is n + 1 long however few distinct values there are. If the interviewer cares more about memory than speed, the heap is the better answer, and being able to make that call is worth more than knowing either one.

What the interviewer is checking

That you read the complexity constraint and understood which solutions it eliminates.
That you do not defend a max-heap with "usually there are few unique values".
If you use a heap, that it is a min-heap capped at k, and you can explain the inversion.
That you spot bounded integer counts as an opening for bucket or counting sort.
Ties at the boundary, and not overrunning the result array inside a bucket.
That you can compare the two good answers on time versus space rather than naming only one.