剑指Offer 41. 数据流中的中位数（Hard）

【题目链接】

题解

1. 数据流中的中位数（优先队列 / 堆，清晰图解）

思路

代码

from heapq import *
'''
heaqp模块提供了堆队列算法的实现，也称为优先级队列算法。
要创建堆，请使用初始化为[]的列表，或者可以通过函数heapify（）将填充列表转换为堆。
提供以下功能：
heapq.heappush（堆，项目）
将值项推入堆中，保持堆不变。
heapq.heapify（x）
在线性时间内将列表x转换为堆。
heapq.heappop（堆）
弹出并返回堆中的最小项，保持堆不变。如果堆是空的，则引发IndexError。
'''

class MedianFinder:
    def __init__(self):
        """
        initialize your data structure here.
        """
        self.A = [] # 小顶堆，存储大的一半元素，A中元素数量可能会比B中元素数量多一个
        self.B = [] # 大顶堆，存储小的一半元素

    # 1.直接使用heappush与heappop的组合操作
    def addNum(self, num: int) -> None:
        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))
    
    # 2.使用heappush来代替heappush与heappop的组合操作，效率更高
    def addNum(self, num: int) -> None:
        # 若A、B中元素数量不等，则应向B中添加一个元素，保持A、B中元素数量一致
        if len(self.A) != len(self.B):
            # Python默认是小顶堆，因此传入和取出相反数则可变相实现大顶堆
            heapq.heappush(self.B, -heapq.heappushpop(self.A, num))
        else:
            heapq.heappush(self.A, -heapq.heappushpop(self.B, -num))

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


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

牛客（含有操作指示）

#
# the medians
# @param operations int整型二维数组 ops
# @return double浮点型一维数组
#
import heapq as hq
class Solution:
    def flowmedian(self , operations ):
        # write code here
        big = []   # 较大的一半的堆（小顶堆）
        small = [] # 较小的一半的堆（大顶堆）
        res = []
        
        for op in operations:
            # 1表示加入中位数
            if op[0] == 1:
                value = op[1]
                # 若较大堆为空 or 当前值val大于较大堆的堆顶元素，则入较大堆
                if len(big) == 0 or value > big[0]:
                    hq.heappush(big, value)
                else:
                    hq.heappush(small, -value)
                    
                # 若一堆比另一堆多2个及以上，则多的一堆出栈元素，再进入少的一堆
                if len(small) - len(big) > 1:
                    hq.heappush(big, -hq.heappop(small))
                elif len(big) - len(small) > 1:
                    hq.heappush(small, -hq.heappop(big))
                    
            # 2表示查询中位数
            elif op[0] == 2:
                n1, n2 = len(small), len(big)
                if n1 == 0 and n2 == 0:
                    res.append(-1)
                # 元素多的一堆的堆顶即为中位数（较小堆需要取反）
                elif n1 > n2:
                    res.append(-small[0])
                elif n2 > n1:
                    res.append(big[0])
                # 若一样多，则两个堆顶取均值
                else:
                    res.append(float(big[0] - small[0])/2)
                    
        return res