代码随想录算法训练营第十三天 | 347. 前 K 个高频元素

347. 前 K 个高频元素

leetcode
文档题解
视频题解

思路：使用 map 存储各个元素出现的频率，然后使用小顶堆存储 K 个元素，这 K 个元素就是高频元素。

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Queue<int[]> queue = new PriorityQueue<>((p1, p2) -> p1[1] - p2[1]);
        Map<Integer, Integer> map = new HashMap<>();

        for (int num: nums) {
            map.put(num, map.getOrDefault(num, 0) + 1);
        }

        for (Map.Entry<Integer, Integer> entry: map.entrySet()) {
            queue.offer(new int[]{ entry.getKey(), entry.getValue() });
            // 当队列的元素数量大于 K，说明满了，要 pop 一些元素，小顶堆 pop 的是当前队列中最小的元素，因此剩下的元素都是较大的，也就是前 K 个频率较高的
            if (queue.size() > k) {
                queue.poll();
            }
        }

        int[] res = new int[k];
        for (int i = 0; i < k; i++) {
            res[i] = queue.poll()[0];
        }

        return res;
    }
}