LeetCode 347. 前 K 个高频元素 | Python

347. 前 K 个高频元素

---

题目来源：力扣（LeetCode）https://leetcode-cn.com/problems/top-k-frequent-elements

题目

---

给定一个非空的整数数组，返回其中出现频率前 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 高的元素。提示中说明，算法时间复杂度要优于 O(nlogn)，又因为只需返回前 k 个频率最高的元素，那么我们借助堆的思路，对于频率 k 之后的不做处理，进而将时间复杂度优化到 O(nlogk)。

那么具体的做法如下：

• 先用哈希表统计元素出现的次数；
• 建立一个最小堆，维护最小堆：
  • 当堆中元素数目小于 k，这里直接将元素插入；
  • 若堆中元素数目等于 k 时，这个时候要将新元素出现频率与堆顶元素出现频率进行比较。若新元素出现频率大于堆顶元素，那么弹出，插入新元素。
• 最终，堆中的 k 个即是要求的答案。

具体代码实现如下（这里直接使用了 heapq 模块）。

from typing import List

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        hash\_map = {}

        # 哈希表统计元素出现频率
        for num in nums:
            if num not in hash\_map:
                hash\_map[num] = 0
            hash\_map[num] += 1
        
        # 建立最小堆，存储频率最大的 k 个元素
        import heapq

        pq = []

        for key in hash\_map:
            if len(pq) < k:
                heapq.heappush(pq, (hash\_map[key],key))
            elif hash\_map[key] > pq[0][0]:
                heapq.heapreplace(pq, (hash\_map[key],key))

        # 取出最小堆中的元素
        res = []
        while pq:
            res.append(pq.pop()[1])

        return res


# nums = [3,0,1,0]
# nums = [4,1,-1,2,-1,2,3]
# k = 2

# solution = Solution()

# print(solution.topKFrequent(nums, k))

---