347. 前 K 个高频元素 - 力扣（LeetCode）

一、题目

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

示例 1：

输入: nums = [1,1,1,2,2,3], k = 2
输出: [1,2]

示例 2：​​​​​​​

输入: nums = [1], k = 1
输出: [1]

提示：

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

二、代码

class Solution {
    // 词频类
    class Node {
        // 数字
        public int num;
        // 该数字在数组中的出现次数
        public int cnt;

        public Node(int num) {
            this.num = num;
            this.cnt = 1;
        }
    }

    // 比较器，用于构建小根堆
    public class CountComparator implements Comparator<Node> {
        @Override
        public int compare(Node o1, Node o2) {
            return o1.cnt - o2.cnt;
        }
    }

    public int[] topKFrequent(int[] nums, int k) {
        // 词频表  key:存放数组中的数字    value:存放数组中数字对应的Node
        HashMap<Integer, Node> countMap = new HashMap<>();
        // 遍历数字，构建词频表
        for (int i = 0; i < nums.length; i++) {
            if (countMap.containsKey(nums[i])) {
                countMap.get(nums[i]).cnt++;
            } else {
                countMap.put(nums[i], new Node(nums[i]));
            }
        }

        // 创建小根堆，根据每个数的出现次数排序，出现次数小的在堆上面
        PriorityQueue<Node> heap = new PriorityQueue<>(new CountComparator());
        // 遍历词频表，将每一个数字的Node加入到小根堆中，维持小根堆的大小不能超过k
        // 当小根堆大小为k，如果遍历到的数字的出现次数比堆顶还要大，那么就是找到了一个更好的数字，我们将小根堆堆顶数字弹出，将这个新的数字加入到堆中
        // 这样遍历完之后我们就可以找到数组中词频数在前K名的所有num了
        for (Node node : countMap.values()) {
            // 堆还没有满或者node.cnt小于此时的堆顶cnt，就将该数的node加入到heap中
            if (heap.size() < k || heap.size() == k && node.cnt > heap.peek().cnt) {
                heap.add(node);
            }
            // 如果上一步是加入了一个比堆顶词频大的Node，这里就要将这个最小的堆顶弹出，维持堆的大小不能超过k
            if (heap.size() > k) {
                heap.poll();
            }
        }
        // 将堆中的数据都加入到答案数组中
        int[] ans = new int[k];
        int index = 0;
        while (!heap.isEmpty()) {
            ans[index++] = heap.poll().num;
        }
        // 返回答案
        return ans;
    }
}

三、解题思路 

先遍历一遍弄出词频表。

小根堆按照出现次数组织，次数少的放在顶部。

如果遍历到的一个新的数，此时堆已经满了，但是新遍历到的数的词频比堆顶的词频还要大，说明这个数更好，我们将堆顶的数弹出，然后将新的数加入到堆中。

相当于我们设置了一个门槛，然后判断当前遍历到的记录能不能把门槛干掉自己进来。

整个过程维持堆的大小不超过k，这样就可以收集到前K个高频元素。