Top K Frequent Elements

mediumArrays & Hashing53 stepsJava · Python · C++

Return the k most frequent values. Count frequencies with a map, then place each value into a bucket indexed by its frequency. Reading buckets from the highest frequency down gives the answer in O(n), no sorting needed.

▶ Watch the animated trace

Traced input: nums = [1, 1, 1, 2, 2, 3], k = 2 → [1, 2]. Original problem: neetcode.io/problems/top-k-elements-in-list

Complexity

BestO(n)

AverageO(n)

WorstO(n)

SpaceO(n)

A frequency can be at most n, so there are n + 1 buckets.

Key idea

Frequencies are bounded by n, so they can be used as array indices.
bucket[f] holds every value that appears exactly f times.
Walk buckets from n down to 1 and collect values until you have k.

Common mistakes

Sorting the map entries by count is O(n log n), fine but not optimal.
A min-heap of size k is O(n log k); know both and when each is better.
Sizing the bucket array as n instead of n + 1 crashes when every element is equal.

Solution

Java

int[] topKFrequent(int[] nums, int k) {
    Map<Integer, Integer> freq = new HashMap<>();
    for (int x : nums) freq.merge(x, 1, Integer::sum);
    List<Integer>[] bucket = new List[nums.length + 1];
    for (Map.Entry<Integer, Integer> e : freq.entrySet()) {
        int f = e.getValue();
        if (bucket[f] == null) bucket[f] = new ArrayList<>();
        bucket[f].add(e.getKey());
    }
    int[] res = new int[k];
    int idx = 0;
    for (int f = bucket.length - 1; f >= 0 && idx < k; f--) {
        if (bucket[f] == null) continue;
        for (int x : bucket[f]) {
            if (idx == k) break;
            res[idx++] = x;
        }
    }
    return res;
}

Python

def topKFrequent(nums: List[int], k: int) -> List[int]:
    freq = Counter(nums)
    bucket = [[] for _ in range(len(nums) + 1)]
    for x, f in freq.items():
        bucket[f].append(x)
    res = []
    for f in range(len(bucket) - 1, 0, -1):
        for x in bucket[f]:
            if len(res) == k:
                break
            res.append(x)
    return res

C++

vector<int> topKFrequent(vector<int>& nums, int k) {
    unordered_map<int, int> freq;
    for (int x : nums) freq[x]++;
    vector<vector<int>> bucket(nums.size() + 1);
    for (auto& [x, f] : freq) {
        bucket[f].push_back(x);
    }
    vector<int> res;
    for (int f = (int)bucket.size() - 1; f > 0 && (int)res.size() < k; f--) {
        for (int x : bucket[f]) {
            if ((int)res.size() == k) break;
            res.push_back(x);
        }
    }
    return res;
}

Edge cases

Each one has its own trace in the interactive view.

Case	Input	Answer
Single element One value in bucket 1.	`nums = [1], k = 1`	`[1]`
All values equal Frequency 3 = n, which is why the array has n + 1 buckets.	`nums = [4, 4, 4], k = 1`	`[4]`
k equals the number of distinct values Every value is returned.	`nums = [1, 2], k = 2`	`[1, 2]`
Negative values Values are map keys, so negatives are fine.	`nums = [-1, -1, 2], k = 1`	`[-1]`
Ties at the cut-off Three values tie; which two you get depends on map order. LeetCode guarantees the answer is unique, so this can't happen there.	`nums = [1, 2, 3, 1, 2, 3], k = 2`	`[1, 2]`