【重点】【堆】347.前K个高频元素

题目

最大的K个元素 => 小根堆（类似上窄下宽的梯形）
最小的K个元素 => 大根堆（类似倒三角形）

法1：小根堆

Python

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        tmp_map = collections.Counter(nums)
        min_heap = []
        for key, val in tmp_map.items():
            if len(min_heap) < k:
                heapq.heappush(min_heap, (val, key))
            else:
                heapq.heappushpop(min_heap, (val, key))
        
        return [item[1] for item in min_heap]

Java

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> valToCountMap = new HashMap<>();
        for (int i : nums) {
            valToCountMap.put(i, valToCountMap.getOrDefault(i, 0) + 1);
        }
        PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> a[1] - b[1]);
        for (Map.Entry<Integer, Integer> entry : valToCountMap.entrySet()) {
            int val = entry.getKey(), count = entry.getValue();
            if (queue.size() == k) {
                if (count > queue.peek()[1]) {
                    queue.poll();
                    queue.offer(new int[]{val, count});
                }
            } else {
                queue.offer(new int[]{val, count});
            }
        }

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

        return res;
    }
}