Top K Frequent Elements

Medium Problem 9 of 151 Java and Python

Practice on LeetCode

Problem

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

The answer is guaranteed to be unique.

Example 1

Input: nums = [1, 1, 1, 2, 2, 3], k = 2
Output: [1, 2]
Why: 1 appears three times and 2 twice.

Example 2

Input: nums = [1], k = 1
Output: [1]

Constraints

1 ≤ nums.length ≤ 10⁵
k is between 1 and the number of distinct elements.
Your solution should be better than O(n log n).

Explanation

First count each value. The question is then "which k counts are largest?"

A count can never be more than n. So instead of sorting, make n + 1 buckets where bucket f holds the values seen exactly f times, and read the buckets from the top down.

Approach 1 · Naive

Count, then sort by count

Time O(n log n)
Space O(n)

Count with a hash map, sort the distinct values by count (highest first) and take the first k.

Count every value.
Sort the distinct values by descending count.
Return the first k.

Java

import java.util.*;

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> count = new HashMap<>();
        for (int num : nums) {
            count.merge(num, 1, Integer::sum);
        }
        List<Integer> values = new ArrayList<>(count.keySet());
        values.sort((a, b) -> count.get(b) - count.get(a));
        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = values.get(i);
        }
        return result;
    }
}

Python

from collections import Counter

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        count = Counter(nums)
        ordered = sorted(count, key=count.get, reverse=True)
        return ordered[:k]

Approach 2 · Optimal

Bucket sort by frequency

Time O(n)
Space O(n)

Put each value in the bucket matching its count, then walk the buckets from the highest count down until you have k values.

Count every value.
Create n + 1 empty buckets; add each value to buckets[count].
Go from bucket n down to 1, collecting values until you have k.

Java

import java.util.*;

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> count = new HashMap<>();
        for (int num : nums) {
            count.merge(num, 1, Integer::sum);
        }
        List<List<Integer>> buckets = new ArrayList<>();
        for (int i = 0; i <= nums.length; i++) {
            buckets.add(new ArrayList<>());
        }
        for (Map.Entry<Integer, Integer> entry : count.entrySet()) {
            buckets.get(entry.getValue()).add(entry.getKey());
        }
        int[] result = new int[k];
        int filled = 0;
        for (int f = nums.length; f > 0 && filled < k; f--) {
            for (int value : buckets.get(f)) {
                if (filled == k) {
                    break;
                }
                result[filled++] = value;
            }
        }
        return result;
    }
}

Python

from collections import Counter

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        count = Counter(nums)
        buckets = [[] for _ in range(len(nums) + 1)]
        for value, freq in count.items():
            buckets[freq].append(value)
        result = []
        for freq in range(len(nums), 0, -1):
            for value in buckets[freq]:
                result.append(value)
                if len(result) == k:
                    return result
        return result

Complexity compared

Approach	Time	Space
Naive Count, then sort by count	O(n log n)	O(n)
Optimal Bucket sort by frequency	O(n)	O(n)

Key takeaways

When the values you sort by are bounded by n, bucket sort gives O(n).
A min-heap of size k (O(n log k)) is the other answer interviewers expect; it also works on a stream.

Your notes

Saved in this browser as you type.

Solved it on your own?

Mark it done to track your progress, or star it to come back before an interview.

Next problem