Top K Frequent Elements [Medium] — Bucket Sort O(n)

Problem Statement

Given an integer array, return the k most frequent elements. O(n log n) is okay, but can you do better?

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

Intuition

Bucket Sort O(n): Create buckets indexed by frequency (0 to n). Place each element in its frequency bucket. Scan from high to low, collecting k elements.

Solutions

C++

vector<int> topKFrequent(vector<int>& nums, int k) {
    unordered_map<int,int> freq;
    for(int n:nums) freq[n]++;
    int n=nums.size();
    vector<vector<int>> buckets(n+1);
    for(auto& [val,cnt]:freq) buckets[cnt].push_back(val);
    vector<int> res;
    for(int i=n;i>=0&&res.size()<k;i--)
        for(int v:buckets[i]) if(res.size()<k) res.push_back(v);
    return res;
}

Java

public int[] topKFrequent(int[] nums, int k) {
    Map<Integer,Integer> freq=new HashMap<>();
    for(int n:nums) freq.merge(n,1,Integer::sum);
    List<Integer>[] buckets=new List[nums.length+1];
    for(Map.Entry<Integer,Integer> e:freq.entrySet()){
        int c=e.getValue();
        if(buckets[c]==null) buckets[c]=new ArrayList<>();
        buckets[c].add(e.getKey());
    }
    int[] res=new int[k]; int idx=0;
    for(int i=nums.length;i>=0&&idx<k;i--)
        if(buckets[i]!=null) for(int v:buckets[i]) if(idx<k) res[idx++]=v;
    return res;
}

Python

from collections import Counter

def topKFrequent(nums: list[int], k: int) -> list[int]:
    freq = Counter(nums)
    buckets = [[] for _ in range(len(nums)+1)]
    for val, cnt in freq.items():
        buckets[cnt].append(val)
    res = []
    for i in range(len(buckets)-1, -1, -1):
        res.extend(buckets[i])
        if len(res) >= k:
            return res[:k]
    return res

# Heap approach O(n log k):
# return [v for v,_ in Counter(nums).most_common(k)]

Complexity

Bucket Sort: O(n) time, O(n) space
Heap: O(n log k) time