[数据结构与算法]优先队列 -leetcode-347.前 K 个高频元素

347. 前 K 个高频元素

题目描述

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

示例 1:

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

解法选择：

• 优先队列
  • 求 top k 问题，优先队列是一把好手

优先队列

• 底层基于堆实现。堆是一种完全二叉树，其父节点大于（或小于）所有子节点，即大根堆（或小根堆）
• 每次增/删后，保持栈顶元素是 max（或 min）的
  • 仅保持局部是有序的，并不是全部有序

注意：

• 到底选用大顶堆还是小顶堆？
  • 一开始我就想也没想，题目要求 top k，那“铁定”大顶堆啊。结果，就是直接翻车 WA，调试发现不对
  • 后来一想如果选大顶堆的话，那么，每次出队列弹出的都是max的元素，必然无法使得保留在优先队列中的最终k个是top k（人家都被弹了，还保留个 der 的 top k）
  • 所以，铁定的选择小顶堆。每次将当前元素 value 和当前队头元素对应的 value 比较，如果大于之，则出队；否则，啥也不管

思路：

1. 使用 map 记录nums中元素出现的个数；
   • key：nums[i]；value：nums[i] 出现的次数
2. 将 map 中前 k 个元素的 key 存入 PriorityQueue；
3. 将 map 中剩余的 size - k 个元素，与队头元素 peek 比较：
   • 如果 map 中当前元素的 value 大于 peek 在 map 中的 value，则将当前元素 key 入队；
   • 否则，不做任何操作。
4. 倒序（因为每次取得都是最小的）从小顶堆中取出k个元素，作为最终结果
   • 但题目，其实没有要求输出顺序（扩展一下而已）

class Solution {
    public int[] topKFrequent(int[] nums, int k) { 
		// 记录nums中元素出现的个数
		Map<Integer, Integer> map = new HashMap<Integer, Integer>();
		for (int i : nums) {
			if (!map.containsKey(i)) {
				map.put(i, 1);
			} else {
				map.put(i, map.get(i) + 1);
			}
		}
		
		// 优先队列（小顶堆k个元素）
		Queue<Integer> queue = new PriorityQueue<>(new Comparator<Integer>() {
			@Override
			public int compare(Integer a, Integer b) {
				return map.get(a) - map.get(b); // // 排序规则（出现次数升序）
			}
		}); 
		
		// 遍历 map，用小顶堆保存频率最大的k个元素
		for (Integer key : map.keySet()) {
			if (queue.size() < k) { // 堆里小于k个元素时候，直接填入
				queue.add(key);
			} else if (map.get(key) > map.get(queue.peek())) { // 核心：保证小根堆里面只有k个元素，并且当前队列元素是最大的
				queue.remove();
				queue.add(key);
			}
		}
		
		// 倒序（因为每次取得都是最小的）从小顶堆中取出k个元素，作为最终结果
        // 但题目，其实没有要求输出顺序（扩展一下而已）
		int cnt = 0; // ans的当前位置
		int[] ans = new int[k]; // 最终答案
		for (int i = k - 1; i >= 0; i--) {
            ans[i] = queue.poll();
        }
		return ans;
    }
}