Leetcode刷题记录——剑指 Offer 41. 数据流中的中位数

我们将数据分开存储
用两个小根堆
一个存数据流偏小的半段
一个存数据流偏大的半段
如何实现呢
python中的heapq是默认的小根堆 也就是每次都弹出最小的元素
那么我们用它来存储偏大的半段 则我们每次都可以吐出偏大一段的最小值

要构造偏小、且能随时吐出自身中最大元素的半段
我们还是用另一个小根堆B代替 只不过在记录数据的时候有一些技巧

我们在每次录入时进行检查
如果小根堆A和大跟堆B长度相等
则我们将元素先取相反数然后录入B
然后再将它吐出后取反录入A

如果二者长度不等
我们先将元素录入A
然后从A中吐出一个取反后录入B

这样 我们永远保证A中元素个数不小于B中元素个数

from heapq import *

class MedianFinder:
    def __init__(self):
        self.A = [] # 小顶堆，保存较大的一半
        self.B = [] # 大顶堆，保存较小的一半

    def addNum(self, num: int) -> None:
        if len(self.A) == len(self.B):
            heappush(self.B,-num)
            heappush(self.A,-heappop(self.B))
        else:
            heappush(self.A,num)
            heappush(self.B,-heappop(self.A))
    def findMedian(self):
        return self.A[0] if len(self.A) != len(self.B) else (self.A[0]-self.B[0])/2