Python实现Leetcode 295 数据流的中位数

原理的话就看东哥的文章吧，这只贴Python代码实现。

import heapq
class MedianFinder:

    def __init__(self):
        self.min_heap = []  # 梯形, min_heap(存大值), 值不小于三角形里面的，维持min_heap元素个数>=max_heap
        self.max_heap = []  # 三角形, max_heap(存小值)，值不大于梯形里面的

    def addNum(self, num: int) -> None:
        if len(self.min_heap) > len(self.max_heap):  # 值最终放入max_heap,三角形
            heapq.heappush(self.min_heap, num)  # 先放入min_heap, 梯形
            tmp = heapq.heappop(self.min_heap)  # 再pop出来
            heapq.heappush(self.max_heap, -tmp)  # 放入 三角形
        else:
            heapq.heappush(self.max_heap, -num)
            tmp = -heapq.heappop(self.max_heap)
            heapq.heappush(self.min_heap, tmp)

    def findMedian(self) -> float:
        if len(self.min_heap) < len(self.max_heap):
            return -self.max_heap[0]  # return的时候不要pop出去了
        elif len(self.min_heap) > len(self.max_heap):
            return self.min_heap[0]
        else:
            return (self.min_heap[0] - self.max_heap[0]) / 2


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

obj = MedianFinder()
obj.addNum(1)
obj.addNum(2)
print(obj.findMedian())  # # 输出1.5
obj.addNum(3)
print(obj.findMedian())  # 输出2