Top K Frequent Elements

NeetCode: Top K Frequent Elements · LeetCode: 347. Top K Frequent Elements · Difficulty: Medium

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

Text

Input:  nums = [1, 2, 2, 3, 3, 3], k = 2    Output: [2, 3]
Input:  nums = [7, 7],             k = 1    Output: [7]

Brute force: count, then sort

Count every value with a HashMap, sort the distinct values by count (highest first), and take the first k. The comparator uses Integer.compare and not b - a, because a subtraction can overflow. HashMap iteration order is unspecified, so values with equal counts may come out in any order; the problem allows any order, and main sorts the result before printing.

Java

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class Main {
    static class Solution {
        public int[] topKFrequent(int[] nums, int k) {
            Map<Integer, Integer> counts = new HashMap<>();
            for (int x : nums) {
                counts.merge(x, 1, Integer::sum);
            }
            List<Map.Entry<Integer, Integer>> items = new ArrayList<>(counts.entrySet());
            items.sort((a, b) -> Integer.compare(b.getValue(), a.getValue())); // highest count first
            int[] result = new int[k];
            for (int i = 0; i < k; i++) {
                result[i] = items.get(i).getKey();
            }
            return result;
        }
    }

    public static void main(String[] args) {
        Solution sol = new Solution();
        int[] a = sol.topKFrequent(new int[] {1, 2, 2, 3, 3, 3}, 2);
        Arrays.sort(a); // the answer may come in any order
        System.out.println(Arrays.toString(a));
        System.out.println(Arrays.toString(sol.topKFrequent(new int[] {7, 7}, 1)));
    }
}

Output

[2, 3]
[7]

Time: O(n log n). Space: O(n).

Optimal: bucket sort by frequency

A value can appear at most n times, so frequencies are small integers. Make n + 1 buckets where bucket f holds the values that appear exactly f times, then read the buckets from highest frequency down until k values are collected. No comparison sort is needed.

Java

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class Main {
    static class Solution {
        public int[] topKFrequent(int[] nums, int k) {
            Map<Integer, Integer> counts = new HashMap<>();
            for (int x : nums) {
                counts.merge(x, 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> e : counts.entrySet()) {
                buckets.get(e.getValue()).add(e.getKey());
            }

            int[] result = new int[k];
            int filled = 0;
            for (int f = buckets.size() - 1; f >= 0 && filled < k; f--) {
                for (int value : buckets.get(f)) {
                    result[filled++] = value;
                    if (filled == k) {
                        break;
                    }
                }
            }
            return result;
        }
    }

    public static void main(String[] args) {
        Solution sol = new Solution();
        int[] a = sol.topKFrequent(new int[] {1, 2, 2, 3, 3, 3}, 2);
        Arrays.sort(a); // the answer may come in any order
        System.out.println(Arrays.toString(a));
        System.out.println(Arrays.toString(sol.topKFrequent(new int[] {7, 7}, 1)));
    }
}

Output

[2, 3]
[7]

Time: O(n). Space: O(n).

A PriorityQueue (a min-heap in Java) limited to k entries also works in O(n log k). Besides the count map it holds only k entries, which helps when k is much smaller than the number of distinct values.