剑指offer41.数据流中的中位数（小顶堆和大顶堆）

题目：

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

例如，

[2,3,4] 的中位数是 3            [2,3] 的中位数是 (2 + 3) / 2 = 2.5

设计一个支持以下两种操作的数据结构：

void addNum(int num) - 从数据流中添加一个整数到数据结构中。
double findMedian() - 返回目前所有元素的中位数。
示例 1：

输入：
["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]
[[],[1],[2],[],[3],[]]
输出：[null,null,null,1.50000,null,2.00000]

示例 2输入：
["MedianFinder","addNum","findMedian","addNum","findMedian"]
[[],[2],[],[3],[]]
输出：[null,null,2.00000,null,2.50000]
 

限制：最多会对 addNum、findMedian 进行 50000 次调用。

解题思路：

建立一个小顶堆A和大顶堆B，各保存列表的一半元素。那么什么叫小顶堆和大顶堆呢？小顶堆和大顶堆，其本质是一颗完全二叉树。除叶子节点外，小顶堆的每个父节点的key都要比其左右两个子节点的key小；大顶堆的每个父节点的key都要比其左右两个子节点的key大。这里说的key暂时理解为节点的取值吧，而index为节点在树中的索引或者位置。小顶堆/大顶堆的特点在于，其根节点一定是整个数中最小或者最大的元素

• A保存较大的一半，长度为N/2（N为偶数）或(N+1)/2（N为奇数）；
• B保存较小的一半，长度为N/2（N为偶数）或(N-1)/2（N为奇数）

然后，中位数可仅根据A,B的堆顶元素计算得到中位数。

算法的流程包括：

设元素总数为N=m+n，其中m和n分别为A和B中的元素个数。

1、addNum(num)

• 当m=n（即N为偶数）：需向A添加一个元素。通过：将新元素num插入B中，再将B堆顶元素插入到A中；
• 当m!=n（即N为奇数）：需向B添加一个元素。通过：将新元素num插入A中，再将A堆顶元素插入到B中

2、findMedian()

• 当m=n（N为偶数）：中位数为（A的堆顶元素+B的堆顶元素）/2
• 当m!=n（N为奇数）：中位数为A的堆顶元素

代码：

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.A, num)
            heappush(self.B, -heappop(self.A))
        else:
            heappush(self.B, -num)
            heappush(self.A, -heappop(self.B))

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

复杂度：

• 时间复杂度：查找中位数为O(1)，获取堆顶元素使用O(1)；添加数字O(logN)，堆的插入和弹出操作使用O(logN)
• 空间复杂度：O(N)，N为数据流中的元素数量，小顶堆和大顶堆中最多同时保存N个元素