优先级队列
- 优先级队列就是一个堆,因为其对外接口是从队头取元素,从队尾添加元素,而没有其他取元素的方式,看起来像一个队列。
- 堆是一个完全二叉树,书中每个节点的值都不小于或者不大于其左右孩子的值,如果父节点的值不小于左右孩子的值,那么这个堆就是大顶堆;父节点的值不大于左右孩子的值,这个堆就是小顶堆;
- 大顶堆堆头是最大元素,小顶堆堆头是最小元素。
- 本题中使用大顶堆还是小顶堆呢?本题中使用小顶堆,若使用大顶堆,因为要保留前k个高频元素,那么每次更新大顶堆的时候,会pop出最大的一个元素,这显然是不能完成任务的,而使用小顶堆就可以每次弹出最小的元素,剩下的就是k个高频元素。
代码
python中heapq提供了对堆的支持,heapq默认是最小堆;
这里主要使用到其两个方法:
heapq.heappush(heap, items)
:heap可以通过列表list来初始化,items是要传入push的值;heapq.heappop(heap)
:删除并返回最小值,也就是pop出heap[0]。
import heapq
class Solution:
def topKFrequent(self, nums: List[int], k: int) -> List[int]:
map_ = {}
for i in range(len(nums)):
map_[nums[i]] = map_.get(nums[i], 0) + 1
pri_que = []
for key, freq in map_.items():
heapq.heappush(pri_que, (freq, key))
if len(pri_que) > k:
heapq.heappop(pri_que)
result = [0] * k
for i in range(k-1, -1, -1):
result[i] = heapq.heappop(pri_que)[1]
return result