优先队列(堆)
- 堆是自带排序功能的数据结构
- 其本质是完全二叉树,但为了方便,用数组存储
- 在使用过程中,保证任意父节点的值一定大于或等于/小于或等于子节点
- Python中用
heapq
库来实现堆 - 使用场景:在查询前n个最大/最小的元素时有很好的性能
优先队列
PriorityQueue
底层实现就是堆heapq
(主要区别是,PriorityQueue
提供的是面对对象的接口,而heapq
只提供用于操作列表的函数)- 当需要求前n个大的元素(只查看不取出),应该使用堆(提供了专门函数
nlargest(n,heap)
和nsmallest(n,heap)
),如LeetCode 1046. 最后一块石头的重量 - 每次取出最大/最小元素,使用优先队列
TopK问题
LeetCode 703. 数据流中的第 K 大元素
设计数据结构,支持向其多次添加元素,并实施返回当前所有元素中第k大的元素ps.这个问题也可以表述为:利用一次遍历,找出数组中第k大的数字
思路:
- 暴力方法是直接用数组模拟,然而每次插入元素需要二分查找,且数组的插入性能差
- 我们需要一个自带排序功能的数据结构——堆
实现:
- 用一个大小限制为k的堆,保存当前已经处理的元素中,前k个最大的元素;
- 每次有新元素,直接插入堆中,(如果堆的大小>k,)再弹出最小的(因此需要小根堆),堆中剩下的最小元素就是所求的[第 K 大元素]
①当前加入的元素比堆中所有元素都小,那么它刚加入堆后就会被弹出,不影响前k个元素的排名
②当前加入的元素大于堆中的一些元素,那么总体的前k大的元素就与之前不同了,我们弹出"前k个"之外的最小元素
from heapq import *
class KthLargest:
"""用一个大小限制为k的堆,保存当前已经处理的元素中,前k个最大的元素
每次有新元素,直接插入堆中,如果堆的大小>k,再弹出最小的,
这样,始终保证堆中的最小元素就是所求的[第 K 大元素]"""
def __init__(self, k: int, nums: List[int]):
self.K = k
self.hp = nums # 保存当前的前k个最大的元素,堆顶为最小,即第k大的元素
heapify(self.hp)
while len(self.hp) > k:
heappop(self.hp)
def add(self, val: int) -> int:
heappush(self.hp, val)
while len(self.hp) > self.K:
heappop(self.hp)
return self.hp[0] # 堆顶元素为前k大的元素中,最小的那个
# Your KthLargest object will be instantiated and called as such:
# obj = KthLargest(k, nums)
# param_1 = obj.add(val)
解法二:快速选择算法
在上面的解法中,我们实际上做了一些不必要的工作:不仅求出了第k大的元素,还一同求出了k-1、k-2大的元素,实际上,TopK问题的最优解法是[快速选择算法]