数据结构与算法——使用优先队列的题目

一、数据流的中位数

此题为leetcode第295题
思路：我们可以维护两个堆，用于存储列表中较小一半的最大堆lo，和用于存储列表中较大一半的最小堆hi。这样可以通过堆的顶部访问输入的中位数。需要满足两个条件：两个堆都是平衡的（或近似平衡）；最大堆包含所有较小的数字，最小堆包含所有较大的数字。最大堆lo允许存储的元素最多比最小堆hi多一个，设我们除了了k个元素，如果k = 2 * n + 1为奇数，则允许lo有n + 1个元素，hi有n个元素；如果k = 2 * n为偶数，则lo和hi都可以有n个元素。

• 添加一个元素num：
  • 将num添加到最大堆lo，因为lo收到了一个元素，所以必须为hi做一个平衡的步骤，从lo中移除最大的元素并将其提供给hi
  • 上一个操作之后，最小堆hi可能会比最大堆lo保留更多的元素，因此从hi中去掉最小的元素给lo

import heapq
class MedianFinder:

    def __init__(self):
        """
        initialize your data structure here.
        """
        self.count = 0
        self.max_heap = []
        self.min_heap = []

    def addNum(self, num: int) -> None:
        self.count += 1
        heapq.heappush(self.max_heap, -num) # python默认小根堆，因此大根堆需要取负数
        max_heap_top = heapq.heappop(self.max_heap)
        heapq.heappush(self.min_heap, -max_heap_top)
        
        if self.count % 2 == 1:
            min_heap_top = heapq.heappop(self.min_heap)
            heapq.heappush(self.max_heap, -min_heap_top)

    def findMedian(self) -> float:
        if self.count % 2 == 1:
            return -self.max_heap[0]
        else:
            return (-self.max_heap[0] + self.min_heap[0]) / 2

二、数据流中的第k大元素

此题为leetcode第703题
思路：建立小根堆，维护长度为k。添加元素时，列表元素小于k时直接添加，否则新元素大于堆顶元素时，替换堆顶元素。

import heapq
class KthLargest:
    def __init__(self, k: int, nums: List[int]):
        self.nums = nums
        self.k = k
        heapq.heapify(self.nums)
        # 留下k个元素，即前k大的
        while len(self.nums) > k:
            heapq.heappop(self.nums)

    def add(self, val: int) -> int:
        if len(self.nums) < self.k:
            heapq.heappush(self.nums, val)
        elif self.nums[0] < val:
            # 新的值更大则更新
            heapq.heapreplace(self.nums, val)

        return self.nums[0]

未完待续