Top K Frequent Elements

1. Problem & Core Objective

The problem

Given an integer array nums and an integer k, return the k most frequent elements. You may return the answer in any order.

Input:  nums = [1,1,1,2,2,3], k = 2      Output: [1, 2]
Input:  nums = [1], k = 1                Output: [1]

Constraints:

• 1 <= nums.length <= 10^5
• -10^4 <= nums[i] <= 10^4
• k is in the range [1, number of distinct elements]
• The answer is guaranteed to be unique (no ambiguous ties at the cut-off)

Follow-up stated in the problem: design an algorithm better than O(n log n).

What the interviewer is actually testing

This is a trade-off question, not a puzzle. There are three good answers and the interview is about choosing between them.

1. Do you know all three approaches? Sort (O(n log n)), heap (O(n log k)), bucket (O(n)). A candidate who gives one has answered; one who gives three and picks has interviewed well.
2. Can you justify a min-heap for finding the largest? This is counter-intuitive and it's the most common follow-up.
3. Do you spot the bound that makes O(n) possible? A frequency can never exceed n. That single observation unlocks bucket sort.
4. The follow-up is written into the problem. "Better than O(n log n)" is an explicit instruction not to stop at sorting.

2. First-Principles Thought Process

Step 1 — Decompose the problem

Two independent sub-problems:

1. Count how often each value appears.
2. Select the k values with the highest counts.

Step 1 is unambiguous — a HashMap<value, count> in one pass, O(n). Nothing to decide.

All the interesting choices are in step 2. Say that explicitly; it structures the rest of the conversation.

Step 2 — Read the constraints

n up to 10^5. Both O(n log n) and O(n) fit comfortably, so raw speed isn't the constraint — the problem's own follow-up is. It asks for better than O(n log n), which means sorting is meant to be your starting point, not your answer.

Step 3 — Selection option A: sort the counts

Take the (value, count) pairs and sort descending by count, then take the first k.

O(d log d) where d is the number of distinct values (≤ n). Simple and correct — but it produces a complete ranking of all d values when you only need the top k. Too much work.

Step 4 — Selection option B: a bounded heap

You don't need everything ranked; you need the top k identified. A heap of size k does exactly that.

The counter-intuitive part: to find the k largest, use a min-heap.

Why: you keep a heap holding the best k seen so far. When a new candidate arrives, you must decide whether it beats the weakest current member — so you want the weakest instantly accessible. In a min-heap of the k largest, the root is the weakest.

O(d log k) — better than O(d log d) when k is small.

Step 5 — Selection option C: exploit the bound

Now the key observation:

A frequency is an integer between 1 and n. It is bounded.

Whenever the thing you'd sort by is a bounded integer, you don't need to compare — you can index by it. Make an array where slot f holds all values appearing exactly f times, then walk it from the top.

That's counting/bucket sort. No comparisons, no log factor: O(n).

Step 6 — Pick

The problem explicitly asks for better than O(n log n), so the bucket approach is the intended answer. But the heap is worth presenting first if k is small, because it uses O(k) space rather than O(n).

3. Solution Paths

Approach 1 — Count, then sort

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

    List<Integer> keys = new ArrayList<>(freq.keySet());
    keys.sort((a, b) -> freq.get(b) - freq.get(a));       // descending by count

    int[] res = new int[k];
    for (int i = 0; i < k; i++) res[i] = keys.get(i);
    return res;
}

• Time: O(n + d log d) — counting plus sorting the distinct values.
• Space: O(d).

Correct, and explicitly rejected by the problem's follow-up. Present it as your baseline, then improve.

(Note: freq.get(b) - freq.get(a) is safe here because counts are small non-negative ints. In general, prefer Integer.compare — see 04.)

Counter-questions on this approach

⭐ "The problem explicitly asks for better than O(n log n). Why are you showing me this?"

As a baseline to improve from — it establishes that the counting half is trivial and all the interesting choices are in the selection half. It's also doing too much work conceptually: it produces a full ranking of all d distinct values when I only need the top k.

"freq.get(b) - freq.get(a) — is that comparator safe?"

Here yes, because counts are small non-negative ints so the subtraction can't overflow. In general it's unsafe and I'd write Integer.compare(freq.get(b), freq.get(a)) — subtraction comparators wrap and silently corrupt the sort when values straddle the int range.

Approach 2 — Min-heap of size k

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

    // min-heap ordered by frequency: the LEAST frequent survivor sits at the root
    PriorityQueue<Integer> heap =
        new PriorityQueue<>(Comparator.comparingInt(freq::get));

    for (int key : freq.keySet()) {
        heap.offer(key);
        if (heap.size() > k) heap.poll();        // evict the weakest
    }

    int[] res = new int[k];
    for (int i = 0; i < k; i++) res[i] = heap.poll();
    return res;
}

How it works. Push every distinct value. Whenever the heap exceeds k, remove the root — which, in a min-heap keyed by frequency, is the least frequent element currently held. After processing everything, exactly the k most frequent remain.

Trace on nums = [1,1,1,2,2,3], k = 2. Frequencies: {1→3, 2→2, 3→1}.

Result: {1, 2} ✓

• Time: O(n + d log k).
• Space: O(d + k).

When this is the best answer: when k is much smaller than d, and especially when the data streams — you can't bucket-sort a stream whose length you don't know, but you can maintain a size-k heap forever.

Counter-questions on this approach

⭐ "You're finding the k largest but using a min-heap. Explain that."

The heap holds the best k seen so far. When a new candidate arrives I need to decide whether it beats the weakest current member — so the weakest is the element I need O(1) access to, because it's the one I evict. A min-heap puts the smallest at the root, which is exactly that element. A max-heap would give me instant access to the best element, which I never need to remove.

⭐ "Your comparator calls freq::get on every comparison. Any risk in that?"

Yes — the map must not change while the heap is in use. If a frequency were mutated mid-flight, the heap's ordering invariant would be violated and poll could return the wrong element. Here the map is fully built before the heap is touched, so it's safe.

"When would you choose this over the bucket approach?"

Two cases. If the input streams — you can't size a bucket array without knowing n. And when k is much smaller than n and memory matters, since this is O(k) space versus bucket sort's O(n).

Approach 3 — Bucket sort by frequency (optimal)

public int[] topKFrequent(int[] nums, int k) {
    // Step 1: count
    Map<Integer, Integer> freq = new HashMap<>();
    for (int x : nums) freq.merge(x, 1, Integer::sum);

    // Step 2: bucket[f] = every value that appears exactly f times
    List<Integer>[] buckets = new List[nums.length + 1];
    for (Map.Entry<Integer, Integer> e : freq.entrySet()) {
        int f = e.getValue();
        if (buckets[f] == null) buckets[f] = new ArrayList<>();
        buckets[f].add(e.getKey());
    }

    // Step 3: walk from the highest frequency down, taking k values
    int[] res = new int[k];
    int idx = 0;
    for (int f = buckets.length - 1; f >= 1 && idx < k; f--) {
        if (buckets[f] == null) continue;
        for (int val : buckets[f]) {
            res[idx++] = val;
            if (idx == k) break;
        }
    }
    return res;
}

How it works. The array is indexed by frequency, not by value. Slot f holds every value occurring f times. Since no value can occur more than n times, n + 1 slots suffice. Reading the array from the end downward visits frequencies in descending order for free — no sorting.

Trace on nums = [1,1,1,2,2,3], k = 2. Frequencies: {1→3, 2→2, 3→1}.

Buckets (index = frequency):

Walking down from index 6: nothing at 6, 5, 4. Index 3 → take 1. Index 2 → take 2. idx == k, stop.

Result: [1, 2] ✓

• Time: O(n) — counting is O(n), filling buckets is O(d) ≤ O(n), and the final scan visits at most n + 1 slots.
• Space: O(n) for the bucket array.

Why it's genuinely O(n) — worth stating, because "sorting in linear time" sounds wrong: it isn't a comparison sort. Comparison sorts have an Ω(n log n) lower bound, but that bound only applies when your only operation is comparing pairs. Here we use the frequency as a direct array index, sidestepping comparison entirely. That's only possible because frequencies are bounded integers.

Counter-questions on this approach

⭐ "You're claiming O(n) sorting. Comparison sorts have an Ω(n log n) lower bound — how do you get around it?"

That bound applies only when comparison is your sole operation on the elements. I never compare frequencies at all — I use the frequency as a direct array index. That's counting sort, and it sidesteps the bound entirely. It's only legal because frequencies are bounded integers: no value can occur more than n times.

⭐ "You allocate n + 1 buckets even when there are three distinct values. Isn't that wasteful?"

It is. This is O(n) space regardless of how small k is — for n = 10^5 with three distinct values that's 100,001 mostly-null references. The heap version is O(k). If memory mattered more than time, the heap is the better engineering choice, and I'd say so rather than calling bucket sort unconditionally optimal.

"Why index by frequency rather than by value?"

Values range to ±10^4 and can be negative, so they aren't valid array indices without offsetting. Frequencies are always positive and bounded by n — which is precisely the property that makes indexing safe.

"new List[n+1] produces an unchecked warning. Why, and does it matter?"

Java forbids generic array creation — generics are erased at runtime, so the array couldn't enforce its element type. The raw form with a warning is standard practice here; the alternative is a List<List<Integer>>, which avoids the issue at the cost of pre-allocating every bucket.

Approach 4 — Quickselect (mention only)

Partition the (value, count) pairs around a pivot until the k-th boundary is in place. O(d) average, O(d²) worst case, O(1) extra space.

Theoretically elegant, fiddly to write correctly under pressure, and no better than bucket sort's guaranteed O(n). Mention it exists; don't write it unless the interviewer asks for O(1) space.

Counter-questions on this approach

⭐ "Quickselect is O(n) average with O(1) extra space — strictly better than bucket sort on memory. Why not use it?"

Its worst case is O(d²), it mutates the input, and it's genuinely fiddly to write correctly under time pressure. Against bucket sort's guaranteed O(n) it offers no time advantage — only the space win. I'd reach for it if O(1) extra space were an explicit requirement.

Comparison

4. Why the Optimal Wins

Against sorting. Sorting produces a full ranking of all d distinct values. You need only the top k — typically a tiny slice. You're computing d log d worth of ordering and discarding nearly all of it.

Against the heap. The heap avoids the full ranking but still pays log k per insertion because it maintains an ordered structure. Bucket sort maintains no order at all — it uses the frequency as a memory address. Array indexing is O(1); heap insertion is O(log k). Removing that factor entirely is the win.

The generalizable principle — say this, it's the transferable part:

When the value you'd sort by is a bounded integer, you can index by it instead of comparing. That turns O(n log n) into O(n).

The same reasoning powers counting sort, radix sort, and the int[26] trick in Valid Anagram.

Honest caveat about the space. Bucket sort allocates n + 1 slots even when only a handful are used — for n = 10^5 with 3 distinct values, that's 100,001 mostly-null references. The heap uses O(k). If memory mattered more than time, the heap would be the better engineering choice. Volunteering that keeps you honest rather than sounding like you're reciting "bucket sort is optimal".

Why O(n) is the floor. Every element must be counted; missing one could change the answer. So O(n) is optimal and this achieves it.

5. Java Prerequisites

Generic array creation

List<Integer>[] buckets = new List[nums.length + 1];       // raw type, compiles with a warning

Java does not allow new List<Integer>[n] — generic array creation is forbidden because generics are erased at runtime and the array couldn't enforce its element type. The idiomatic workarounds:

List<Integer>[] buckets = new List[n + 1];                  // unchecked warning; standard practice
@SuppressWarnings("unchecked")
List<Integer>[] b = (List<Integer>[]) new List[n + 1];      // silence it explicitly

List<List<Integer>> buckets = new ArrayList<>();            // or avoid arrays entirely
for (int i = 0; i <= n; i++) buckets.add(new ArrayList<>());

Mention the warning if you use the array form — interviewers notice whether you know why it's there.

Lazy bucket initialization

if (buckets[f] == null) buckets[f] = new ArrayList<>();

new List[n+1] fills with null, not empty lists. Allocating all n + 1 lists up front would waste memory when only a few are used — so create on demand, and null-check when reading.

Heap with a comparator over an external map

PriorityQueue<Integer> heap = new PriorityQueue<>(Comparator.comparingInt(freq::get));

The heap stores values, but orders them by their frequency, looked up in the map. freq::get is the key extractor.

One caveat: the comparator reads the map on every comparison, so the map must not change while the heap is in use — mutating a frequency would corrupt the heap's invariant.

Map.merge for counting

freq.merge(x, 1, Integer::sum);

"If x is absent store 1; if present, combine the old value with 1 by addition." See 02 §1.3.

Min-heap vs max-heap

new PriorityQueue<>();                                 // min-heap — smallest at the root
new PriorityQueue<>(Collections.reverseOrder());       // max-heap
new PriorityQueue<>(Comparator.comparingInt(freq::get));            // min by frequency
new PriorityQueue<>(Comparator.comparingInt(freq::get).reversed()); // max by frequency

For the size-k pattern you want the min-heap, so the weakest survivor is at the root, ready to evict.

6. Interview Communication Guide

Clarifying questions

1. "Is k guaranteed to be at most the number of distinct elements?" — the constraints say yes; confirming avoids defensive code.
2. "What if there are ties at the boundary — is any valid answer acceptable?" — the problem guarantees uniqueness, which removes an entire class of complexity.
3. "Does the output order matter?" — no, per the problem. Don't sort the result.
4. "Is this a one-off array, or a stream?" — a stream rules out bucket sort and makes the heap the answer.
5. "Any memory constraint?" — bucket sort allocates O(n) regardless of k.

The pitch

"This splits into two parts: count the frequencies, then select the top k. Counting is a one-pass HashMap, O(n) — no decisions there. All the choices are in the selection.

The simplest selection is to sort the distinct values by count, O(d log d). But that produces a full ranking when I only need the top k, and the problem explicitly asks for better than O(n log n).

Better: a min-heap of size k. I push each distinct value and evict the root whenever the heap exceeds k. It's a min-heap because I need instant access to the weakest of my current best k — that's the one to drop. O(d log k).

Best: notice a frequency can never exceed n. So instead of sorting by frequency, I can index by it — an array where slot f holds all values appearing f times. Walking it from the top down gives descending frequency order for free. That's O(n), no comparisons.

I'll go with bucket sort. The trade-off is that it allocates n + 1 slots regardless of k, so if memory were tight I'd use the heap instead."

Edge cases to raise proactively

The negative-values point is worth volunteering: "Values range to −10⁴, so I can't index the bucket array by value — but I'm indexing by frequency*, which is always positive, so it's fine."* It shows you checked that the bound you're relying on is the right one.

Bucket index 0 is never used — a value in the map appeared at least once. Starting the downward scan at f >= 1 makes that explicit.

7. Follow-Up Questions — Modified Constraints

The interviewer changes a constraint of the original problem and asks you to solve it again. These are new problems, asked after your solution is accepted — not challenges to it. (Those are the counter-questions attached to each approach in §3.) ⭐ marks the most likely.

⭐ "What if the input is a stream of unknown length?"

Bucket sort is out — you can't size the array without knowing n. The size-k min-heap works indefinitely: O(log k) per element, O(k) memory, and the answer is always available.

"Can you do it in O(1) extra space?"

Quickselect over the (value, count) pairs — partition until the k-th boundary lands in place. O(d) average, O(d²) worst case. It mutates the array and is easy to get wrong under pressure, so I'd only reach for it if space were the binding constraint.

⭐ "Top k frequent words, with ties broken alphabetically?"

Bucket sort no longer suffices, because it gives no order within a bucket. Either sort each bucket as you consume it, or use a heap with a two-key comparator: frequency ascending, then word descending — so the "worst" element by both criteria sits at the root for eviction. That's the multi-key comparator pattern.

"What if the array is enormous and distributed?"

Count locally on each machine, then merge the partial counts by key — the classic MapReduce word-count. Each machine can send only its local top k, since a global top-k element must be top-k somewhere... which is not actually true in general, so the safe version merges full counts. Worth flagging the subtlety rather than asserting the shortcut.