Top K Frequent Elements — Heap vs Bucket Sort

Top K Frequent Elements (Medium) shows up in product analytics and is a classic “two valid approaches” interview question. You can solve it with a min-heap in O(n log k) or with bucket sort in O(n).

The Problem

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

Example:

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

Intuition

The brute force is to count frequencies and sort them by count. That is O(n log n). We can do better.

There are two clean optimal approaches. The heap approach maintains a min-heap of size k over (frequency, value) pairs. When the heap grows past k, pop the smallest. At the end the heap holds the top k. That is O(n log k) time.

The bucket sort approach exploits a key observation: frequency is bounded by len(nums). So we can create an array of buckets where index = frequency, and walk it from high to low collecting values until we have k. That is O(n) time. For small k, the heap is great. For very large k, bucket sort wins.

Explanation with Example

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

• Counts: 1 -> 3, 2 -> 2, 3 -> 1.
• Buckets index by frequency: index 1 has [3], index 2 has [2], index 3 has [1].
• Walk from highest frequency down: pick 1 then 2. Stop at k.

Result: [1, 2].

nums = [1,1,1,2,2,3]    n = 6

counts:  {1:3, 2:2, 3:1}

buckets (index = freq):
index:  0    1    2    3    4    5    6
       [ ] [ 3 ] [ 2 ] [ 1 ] [ ] [ ] [ ]

walk freq high -> low, collect until k=2:
freq 3 -> take 1
freq 2 -> take 2     done
result = [1, 2]

push (3,1)   heap: [(3,1)]
push (2,2)   heap: [(2,2),(3,1)]                size=2=k
push (1,3)   heap: [(1,3),(3,1),(2,2)]
           size > k -> pop smallest (1,3)
           heap: [(2,2),(3,1)]

final heap holds the top k by frequency

Code

from collections import Counter

def topKFrequent(nums, k):
    counts = Counter(nums)
    buckets = [[] for _ in range(len(nums) + 1)]
    for num, freq in counts.items():
        buckets[freq].append(num)
    result = []
    for freq in range(len(buckets) - 1, 0, -1):
        for num in buckets[freq]:
            result.append(num)
            if len(result) == k:
                return result
    return result

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> counts = new HashMap<>();
        for (int x : nums) counts.merge(x, 1, Integer::sum);
        List<List<Integer>> buckets = new ArrayList<>();
        for (int i = 0; i <= nums.length; i++) buckets.add(new ArrayList<>());
        for (var e : counts.entrySet()) buckets.get(e.getValue()).add(e.getKey());
        int[] result = new int[k];
        int idx = 0;
        for (int f = buckets.size() - 1; f > 0 && idx < k; f--) {
            for (int v : buckets.get(f)) {
                if (idx == k) break;
                result[idx++] = v;
            }
        }
        return result;
    }
}

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int,int> counts;
        for (int x : nums) counts[x]++;
        vector<vector<int>> buckets(nums.size() + 1);
        for (auto& p : counts) buckets[p.second].push_back(p.first);
        vector<int> result;
        for (int f = (int)buckets.size() - 1; f > 0 && (int)result.size() < k; f--) {
            for (int v : buckets[f]) {
                result.push_back(v);
                if ((int)result.size() == k) return result;
            }
        }
        return result;
    }
};

Edge Cases

• k equals number of distinct elements: return all of them.
• Empty array with k = 0: return an empty list.
• All elements equal: one bucket, return that single element.
• Negative numbers: counters and heaps handle them identically.
• Ties: the problem allows any valid ordering; both implementations are fine.

Complexity Analysis

• Sort-based brute force: O(n log n) time, O(n) space.
• Heap: O(n log k) time, O(n + k) space.
• Bucket sort: O(n) time, O(n) space.

For small k, heap is great. For very large k, bucket sort wins.

How to Explain It in an Interview

Start by noting that frequencies are integers between 1 and n. That observation unlocks bucket sort. Walk through both. State the trade-off: heap is more general because it doesn’t depend on the bound on frequency; bucket sort is faster but assumes that bound. Then code the one you’re asked for, narrating invariants like “the heap always holds the current top-k candidates by frequency.”

Related Problems

• Top K Frequent Words
• Kth Largest Element in an Array
• Sort Characters By Frequency
• K Closest Points to Origin

Wrap up

Top K Frequent Elements is one of those problems where the optimal solution depends on input structure. Knowing that frequency is bounded by n is the kind of observation that separates senior candidates.