Topics › Heaps › Hard Problems

Top K Frequent Elements

Medium TUF ID 1018 A2ZSDEBlind 75

Problem

Given an integer array nums and an integer k, return the k values that occur most often. The answer may be in any order. The input guarantees that the answer is unique: there is never a tie for the last place in the top k.

Constraints: 1 ≤ nums.length ≤ 10^5, -10^4 ≤ nums[i] ≤ 10^4, and k is between 1 and the number of distinct values. The follow-up asks for something better than O(n log n).

Example 1

nums = [1, 1, 1, 2, 2, 3], k = 2
answer = [1, 2]          // 1 occurs 3 times, 2 occurs twice, 3 once

Example 2

nums = [4, 4, -1, -1, -1, 2], k = 1
answer = [-1]            // -1 occurs 3 times, more than any other value

Every approach starts the same way: one pass with a HashMap from value to count. Write d for the number of distinct values. The approaches differ in how they pick the k largest counts out of those d entries.

Approach 1: Brute force, count and sort

Intuition. Once we know every value’s count, “the k most frequent” is “the first k after sorting by count, descending”. So put the map’s entries into a list, sort it with a comparator on the count, and take the first k keys. It is short and correct, and it does more work than necessary, because it fully orders all d values when we only need the top k.

Algorithm.

  1. Count occurrences with freq.merge(x, 1, Integer::sum).
  2. Copy freq.entrySet() into a list and sort it by value, descending.
  3. Return the keys of the first k entries.
static int[] bruteForce(int[] nums, int k) {
    Map<Integer, Integer> freq = new HashMap<>();
    for (int x : nums) freq.merge(x, 1, Integer::sum);
    List<Map.Entry<Integer, Integer>> entries = new ArrayList<>(freq.entrySet());
    entries.sort((p, q) -> Integer.compare(q.getValue(), p.getValue()));   // most frequent first
    int[] res = new int[k];
    for (int i = 0; i < k; i++) res[i] = entries.get(i).getKey();
    return res;
}

Complexity. Time O(n + d log d): one counting pass plus sorting d entries. In the worst case d = n, so O(n log n). Space O(d).

Approach 2: Better, size-k min-heap

Intuition. This is the standard “top k” pattern. Keep a heap of at most k candidates, ordered by count, with the least frequent candidate on top. Offer each distinct value. If the heap grows to k + 1, evict the top, which is the weakest candidate. A value that is evicted had at least k values with a higher or equal count, so it cannot be in the answer. When every value has been offered, the heap holds exactly the top k.

The heap stores values, and the comparator looks their counts up in freq. The map is never reassigned, so the lambda can capture it.

Algorithm.

  1. Count occurrences.
  2. For each distinct value x: push x, and if the heap size exceeds k, poll.
  3. Poll the remaining k values into the result. Filling from the back gives the most frequent first, though any order is accepted.
static int[] better(int[] nums, int k) {
    Map<Integer, Integer> freq = new HashMap<>();
    for (int x : nums) freq.merge(x, 1, Integer::sum);
    PriorityQueue<Integer> pq = new PriorityQueue<>((p, q) -> Integer.compare(freq.get(p), freq.get(q)));   // least frequent on top
    for (int x : freq.keySet()) {
        pq.add(x);
        if (pq.size() > k) pq.poll();                      // evict the weakest candidate
    }
    int[] res = new int[k];
    for (int i = k - 1; i >= 0; i--) res[i] = pq.poll();
    return res;
}

Complexity. Time O(n + d log k): each of the d offers costs O(log k) because the heap never holds more than k + 1 items. Space O(d + k).

When k is small, this is effectively linear. When k is close to d, it is no better than sorting. The follow-up wants an approach that does not depend on k.

Approach 3: Optimal, bucket sort by frequency

Intuition. Counts are not arbitrary numbers. Every count is an integer between 1 and n. When the keys being sorted come from a small known range, comparison sorting is unnecessary: make one bucket per possible count, drop each value into the bucket for its count, then read the buckets from the highest count down until k values have been collected.

buckets.get(f) holds every value that occurs exactly f times. Walking f from n down to 1 visits values in non-increasing order of frequency, which is exactly the order we want.

Algorithm.

  1. Count occurrences.
  2. Create n + 1 empty lists, indexed by count 0..n.
  3. For each entry (value, count), append value to buckets.get(count).
  4. For f from n down to 1, and while fewer than k values have been collected, copy the values of buckets.get(f) into the result.
static int[] optimal(int[] nums, int k) {
    Map<Integer, Integer> freq = new HashMap<>();
    for (int x : nums) freq.merge(x, 1, Integer::sum);
    List<List<Integer>> buckets = new ArrayList<>();       // buckets.get(f) = values seen exactly f times
    for (int f = 0; f <= nums.length; f++) buckets.add(new ArrayList<>());
    for (Map.Entry<Integer, Integer> e : freq.entrySet()) buckets.get(e.getValue()).add(e.getKey());
    int[] res = new int[k];
    int idx = 0;
    for (int f = nums.length; f >= 1 && idx < k; f--) {
        for (int x : buckets.get(f)) {
            if (idx == k) break;
            res[idx++] = x;
        }
    }
    return res;
}

Dry run on example 1, n = 6, counts {1: 3, 2: 2, 3: 1}:

f bucket collected
6, 5, 4 empty  
3 [1] [1]
2 [2] [1, 2], which is k values, so stop

Complexity. Time O(n): counting is O(n), creating the buckets is O(n), and the final walk touches each bucket and each distinct value at most once. Space O(n) for the buckets.

Another O(n)-average option is quickselect on the distinct values by count, as in “k-th largest element”. It uses less memory than the buckets, but it is harder to get right and quadratic in the worst case. The bucket approach is the one usually expected.

Edge cases and pitfalls