Amazon — Top K Frequent Elements (Bucket Sort / Heap)

Sanjeev SharmaSanjeev Sharma
2 min read

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

ApproachTimeSpace
Bucket sortO(n)O(n)
Heap nlargestO(n log k)O(n)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro