算法学习笔记——数据结构：优先队列/堆（TopK问题）

优先队列（堆）

• 堆是自带排序功能的数据结构
• 其本质是完全二叉树，但为了方便，用数组存储
• 在使用过程中，保证任意父节点的值一定大于或等于/小于或等于子节点
• 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问题的最优解法是[快速选择算法]