leetcode刷题（44）——347.前 k 个高频元素

一、题目

给定一个非空的整数数组，返回其中出现频率前 k 高的元素。

示例 1:

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

示例 2:

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

提示：

• 你可以假设给定的 k 总是合理的，且 1 ≤ k ≤ 数组中不相同的元素的个数。
• 你的算法的时间复杂度必须优于 O(n log n) , n 是数组的大小。
• 题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的。
• 你可以按任意顺序返回答案。

二、思路

要找出出现频率前 k 高的元素，就需要一个能够根据出现频率快速获取元素的数据结构，我们可以使用优先级队列（堆）来完成：

• 使用哈希表来建立数字和其出现次数的映射，遍历一遍数组统计元素的频率；
• 维护一个元素数目为 k 的最小堆；
• 遍历 map，将元素存入堆中。每次都将新的元素与堆顶元素（堆中出现频率最小的元素）进行比较，如果新的元素的频率比堆顶元素的大，则弹出堆顶的元素，将新的元素存入堆中；
• 最终，堆中的 k 个元素即为前 k 个高频元素。

三、代码

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        // 使用map统计数组中各元素出现的次数
        Map<Integer, Integer> map = new HashMap<>();
        for(int i = 0; i < nums.length; i++){
            if(!map.containsKey(nums[i]))
                map.put(nums[i], 1);
            else
                map.put(nums[i], map.get(nums[i]) + 1);
        }
        // 创建最小堆
        PriorityQueue<Integer> heap = 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(heap.size() < k)
                heap.add(key);
            else if(map.get(key) > map.get(heap.peek())){
                heap.remove();
                heap.add(key);
            }
        }
        // 取出最小堆中的元素
        int[] res = new int[k];
        for(int i = 0; i < k; i++){
            res[i] = heap.remove();
        }
        return res;
    }
}