【LeetCode笔记】347. 前K个高频元素（Java、优先队列、小顶堆、HashMap）

题目描述

• 时间复杂度小于O(n*logn)，否则直接sort，再遍历就很轻松。
• 很有学习价值的题目，第一次使用了优先队列PriorityQueue。
  

思路 & 代码

• 首先遍历数组，用HashMap存储 元素值 - 频率 映射
• 然后维护一个容量为 k 的小顶堆（重写比较器，注意写法），遍历哈希表存入小顶堆中
• 接着把最终的小顶堆存入数组即可

class Solution {
    // TreeMap重写comparator？
    public int[] topKFrequent(int[] nums, int k) {
        // 先用HashMap存储 元素 - 频率 映射 O(n) O(n)
        Map<Integer, Integer> hashmap = new HashMap();
        for(int i = 0; i < nums.length; i++){
            if(hashmap.containsKey(nums[i])){
                hashmap.put(nums[i], hashmap.get(nums[i]) + 1);
            }
            else{
                hashmap.put(nums[i], 1);
            }
        }
        // PriorityQueue, 维护一个容量为k的优先队列：小顶堆
        PriorityQueue<Integer> pq = new PriorityQueue<>(new Comparator<Integer>(){
            // 设定以频率分高下
            @Override
            public int compare(Integer key1, Integer key2){
                return hashmap.get(key1) - hashmap.get(key2);
            }
        });
        // 遍历哈希表，轮流对比存入优先队列中 O(n)
        for(Integer key : hashmap.keySet()){
            // 未满情况
            if(pq.size() < k){
                pq.add(key);
            }
            // 满了情况，对比替换
            else if(hashmap.get(key) > hashmap.get(pq.peek())){
                // 移头，新增
                pq.remove();
                pq.add(key);
            }
        }

        // 获取答案 O(n)
        int[] ans = new int[k];
        for(int i = 0; i < k; i++){
            ans[i] = pq.remove();
        }
        return ans;
    }
}

更新版：引入 stream 流 + Lambda

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();
        for(int temp : nums) {
            if(map.containsKey(temp)) {
                map.put(temp, map.get(temp) + 1);
            }
            else {
                map.put(temp, 1);
            }
        }

        PriorityQueue<Integer> queue = new PriorityQueue<>(
            (Integer a, Integer b) -> (map.get(a) - map.get(b)));
        for(Map.Entry<Integer, Integer> entry : map.entrySet()) {
            if(queue.size() < k) {
                queue.add(entry.getKey());
            }
            else if(entry.getValue() > map.get(queue.peek())) {
                queue.remove();
                queue.add(entry.getKey());
            }
        }

        return queue.stream().mapToInt(Integer::valueOf).toArray();
    }
}