347. 前 K 个高频元素

Difficulty
Medium
Tags
堆
URL
https://leetcode-cn.com/problems/top-k-frequent-elements/
Star
给你一个整数数组 nums 和一个整数 k ，请你返回其中出现频率前 k 高的元素。你可以按 任意顺序 返回答案。
示例 1:
输入:nums = [1,1,1,2,2,3], k = 2
输出:[1,2]
示例 2:
输入:nums = [1], k = 1
输出:[1]
提示：
  • 1 <= nums.length <= 105
  • k 的取值范围是 [1, 数组中不相同的元素的个数]
  • 题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的
进阶：你所设计算法的时间复杂度 必须 优于 O(n log n) ，其中 n 是数组大小。

堆

思路
最朴素的思路就是统计所有数字的频率，然后排序得到前k 个高频元素。这样时间复杂度为nlog(n)。
还有一种思路是使用堆，且最好使用小根堆，因为使用大根堆需要将所有数据都放进去的，空间负载杂多较高，而使用小根堆只需要O(k)的空间复杂度。
即维护大小为K的小根堆，如果新来的元素频率小于堆顶，那么跳过，否则加进来。
谨记：前 k 大，用小根堆，求前 k 小，用大根堆。
题解
import heapq
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        cnt = Counter(nums)
        pair = [[v, k] for k, v in cnt.items()]
        
        h = []
        for p in pair:            
            while len(h) > k:
                heapq.heappop(h)

            if len(h) == k and p[0] < h[0][0]:
                continue

            heapq.heappush(h, p)
        
        res = [p[1] for p in h]
        return res