Top K Frequent Elements

1package pkg5leetcode;2 3/*4 * LeetCode 347: Top K Frequent Elements  (Medium)5 * ------------------------------------------------6 * Return the k most frequent elements.7 *8 * APPROACH: count frequencies, then bucket sort by frequency (index = count).9 * COMPLEXITY: Time O(n), Space O(n).  (A heap gives O(n log k).)10 */11import java.util.*;12 13public class leetcode12TopKFrequent {14 15    @SuppressWarnings("unchecked")16    static int[] topKFrequent(int[] nums, int k) {17        Map<Integer, Integer> freq = new HashMap<>();18        for (int n : nums) freq.merge(n, 1, Integer::sum);19 20        // buckets[i] = list of numbers that appear i times21        List<Integer>[] buckets = new List[nums.length + 1];22        for (var e : freq.entrySet()) {23            int c = e.getValue();24            if (buckets[c] == null) buckets[c] = new ArrayList<>();25            buckets[c].add(e.getKey());26        }27 28        int[] res = new int[k];29        int idx = 0;30        for (int c = buckets.length - 1; c >= 0 && idx < k; c--) {31            if (buckets[c] == null) continue;32            for (int num : buckets[c]) {33                if (idx == k) break;34                res[idx++] = num;35            }36        }37        return res;38    }39 40    public static void main(String[] args) {41        int[] r = topKFrequent(new int[]{1, 1, 1, 2, 2, 3}, 2);42        Arrays.sort(r);43        check(Arrays.equals(r, new int[]{1, 2}), "k=2");44        check(Arrays.equals(topKFrequent(new int[]{1}, 1), new int[]{1}), "single");45        System.out.println("leetcode12TopKFrequent: all tests passed");46    }47 48    static void check(boolean cond, String name) {49        if (!cond) throw new AssertionError("FAILED: " + name);50        System.out.println("  PASS " + name);51    }52}