剑指 Offer 41. 数据流中的中位数（大根堆与小根堆）

题目描述

如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。

思路分析

这道题是不定长数据，求其中中位数。

这题直接暴力思想，每次加入一个数都存到res数组里，要求中位数的时候，直接排序一遍res，然后取中间值就行了， 不过这样时间复杂度很高，超时。。。。

正经思路：
利用两个堆来做。

• 堆A，小根堆，总是保存所有数中较大的那一半
• 堆B，大根堆，总是保存所有数中较小的那一半

步骤：
addNum（）函数
进来一个值

• 如果此时 堆A和堆B的长度相等，即len(A) == len(B)，则该值可能属于大的一部分（A堆），也可能属于小的一部分（B堆），所以要从小到大筛选，操作方法：先让该值进入B堆，然后让B堆的堆顶放入A堆。
• 若此时 len(A)!=len(B),则必然是A堆多了一个值，具体为啥可以想一遍流程就知道了。那么此时需要往B堆里加值，操作方法：先让该值进入A堆，然后让A堆顶进入B堆。

findMedian（）函数
这个就简单了。
梳理一下刚才建立的A和B堆。
A为小根堆，存的都是较大的那一半值，所以A堆顶就是 较大值那一半里的最小值。
B为大根堆，同理，B堆顶是较小值那一半里的最大值。

所以中位数：

• 若A堆和B堆数量一样，则中位数=(A堆顶值+B堆顶值)/2
• 数量不一样，则为A堆顶值。

python的heapq里只有小根堆的方法，大根堆的话，直接在操作前面加上负号就行了。

完整代码

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

    def addNum(self, num: int) -> None:
        # 若A和B相等，则A优先
        # 若A和B不相等，则A必然多一个数字，B优先
        
        if len(self.A)!=len(self.B): 
            heappush(self.A,num)
            heappush(self.B,-heappop(self.A))
        else:
            heappush(self.B,-num)
            heappush(self.A,-heappop(self.B))


    def findMedian(self) -> float:
        if len(self.A) == len(self.B):
            return (self.A[0]-self.B[0])/2.0
        else:
            return self.A[0]

# Your MedianFinder object will be instantiated and called as such:
# obj = MedianFinder()
# obj.addNum(num)
# param_2 = obj.findMedian()