Amazon — Top K Frequent Elements (Bucket Sort / Heap)
Advertisement
Problem (Amazon Frequency Classic)
Given array nums and integer k, return the k most frequent elements.
Example:
nums = [1,1,1,2,2,3], k = 2 → [1, 2]
Solutions
Python — Bucket Sort O(n)
from collections import Counter
def topKFrequent(nums, k):
freq = Counter(nums)
buckets = [[] for _ in range(len(nums) + 1)]
for num, cnt in freq.items():
buckets[cnt].append(num)
res = []
for i in range(len(buckets) - 1, 0, -1):
for num in buckets[i]:
res.append(num)
if len(res) == k:
return res
return res
Python — Heap O(n log k)
from collections import Counter
import heapq
def topKFrequent(nums, k):
return heapq.nlargest(k, Counter(nums).keys(), key=Counter(nums).get)
JavaScript
function topKFrequent(nums, k) {
const freq = new Map();
for (const n of nums) freq.set(n, (freq.get(n)||0)+1);
const buckets = Array.from({length:nums.length+1},()=>[]);
for (const [n,f] of freq) buckets[f].push(n);
const res=[];
for (let i=buckets.length-1; i>=0&&res.length<k; i--) res.push(...buckets[i]);
return res.slice(0,k);
}
Java
import java.util.*;
public int[] topKFrequent(int[] nums, int k) {
Map<Integer,Integer> f=new HashMap<>(); for(int n:nums) f.merge(n,1,Integer::sum);
List<Integer>[] b=new List[nums.length+1];
for(int i=0;i<=nums.length;i++) b[i]=new ArrayList<>();
for(Map.Entry<Integer,Integer> e:f.entrySet()) b[e.getValue()].add(e.getKey());
int[] res=new int[k]; int idx=0;
for(int i=nums.length;i>=0&&idx<k;i--) for(int n:b[i]) if(idx<k) res[idx++]=n;
return res;
}
C++
#include <unordered_map>
#include <vector>
using namespace std;
vector<int> topKFrequent(vector<int>& nums, int k) {
unordered_map<int,int> f; for(int n:nums) f[n]++;
vector<vector<int>> b(nums.size()+1);
for(auto& [n,cnt]:f) b[cnt].push_back(n);
vector<int> res;
for(int i=b.size()-1;i>=0&&(int)res.size()<k;i--) for(int n:b[i]) if((int)res.size()<k) res.push_back(n);
return res;
}
C
/* C: qsort by frequency, return top k */
Complexity
| Approach | Time | Space |
|---|---|---|
| Bucket sort | O(n) | O(n) |
| Heap nlargest | O(n log k) | O(n) |
Advertisement