LeetCodeHot100 347. 前 K 个高频元素

题目

给你一个整数数组 nums 和一个整数 k ，请你返回其中出现频率前 k 高的元素。你可以按 任意顺序 返回答案。

示例 1：

**输入：**nums = [1,1,1,2,2,3], k = 2

输出：[1,2]

示例 2：

**输入：**nums = [1], k = 1

输出：[1]

示例 3：

**输入：**nums = [1,2,1,2,1,2,3,1,3,2], k = 2

输出：[1,2]

提示：

1 <= nums.length <= 105
-104 <= nums[i] <= 104
k 的取值范围是 [1, 数组中不相同的元素的个数]
题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的

进阶： 你所设计算法的时间复杂度必须优于 O(n log n) ，其中 n是数组大小。

思路解析

思路很简单，维护一个最小堆，这个堆大小为k，以元素的出现次数排序，达到k后看堆顶的元素次数进行对比然后判断是否移出去。

详细代码与注释

java 复制代码

class Solution {

    public int[] topKFrequent(int[] nums, int k) {

        /* =====================================================
         * 一、统计每个数字出现的次数
         * ===================================================== */

        // HashMap：key 是数字本身，value 是该数字出现的次数
        Map<Integer, Integer> count = new HashMap<>();

        // 遍历数组 nums
        for (int num : nums) {
            /*
             * count.getOrDefault(num, 0)
             * 如果 map 中已经存在 num，则返回 num 对应的次数
             * 如果 map 中不存在 num，则返回默认值 0
             */
            count.put(num, count.getOrDefault(num, 0) + 1);
        }

        /* =====================================================
         * 二、创建一个小顶堆（按出现次数排序）
         * ===================================================== */

        /*
         * PriorityQueue<int[]>：
         * 堆中存放的是 int 数组
         * 约定数组结构：
         *   arr[0] = 数字
         *   arr[1] = 该数字出现的次数
         */
        PriorityQueue<int[]> queue = new PriorityQueue<>(
            new Comparator<int[]>() {
                @Override
                public int compare(int[] o1, int[] o2) {
                    /*
                     * 比较规则：按出现次数升序排序
                     * o1[1] - o2[1] < 0  → o1 优先
                     * 因此这是一个「小顶堆」
                     */
                    return o1[1] - o2[1];
                }
            }
        );

        /* =====================================================
         * 三、遍历 HashMap，把前 k 个高频元素放入堆中
         * ===================================================== */

        // 遍历 Map 中的每一对 (数字, 出现次数)
        for (Map.Entry<Integer, Integer> entry : count.entrySet()) {

            // 当前数字
            int num = entry.getKey();

            // 当前数字出现的次数
            int freq = entry.getValue();

            // 如果堆的大小已经等于 k，说明堆满了
            if (queue.size() == k) {

                /*
                 * peek() 查看堆顶元素（不删除）
                 * 由于是小顶堆，堆顶元素的出现次数是当前堆中最小的
                 */
                if (queue.peek()[1] < freq) {
                    // 当前数字出现次数更大，应该进入前 k

                    // 删除堆顶（出现次数最小的那个）
                    queue.poll();

                    // 把当前数字及其出现次数加入堆中
                    queue.offer(new int[]{num, freq});
                }

                // 如果当前频率 <= 堆顶频率
                // 说明排不进前 k，直接跳过

            } else {
                // 如果堆还没满 k 个元素，直接加入
                queue.offer(new int[]{num, freq});
            }
        }

        /* =====================================================
         * 四、从堆中取出结果
         * ===================================================== */

        // 结果数组，长度为 k
        int[] res = new int[k];

        // 依次从堆中取出元素
        for (int i = 0; i < k; i++) {
            /*
             * poll() 会取出并删除堆顶元素
             * 堆中存的是 [数字, 出现次数]
             * 我们只需要数字，所以取下标 0
             */
            res[i] = queue.poll()[0];
        }

        // 返回前 k 个高频元素
        return res;
    }
}