题解 | #数据流中的中位数#

# -*- coding:utf-8 -*-
import heapq

class Solution:
    def __init__(self) -> None:
        self.min_heap = []
        self.max_heap = []
		# 维护两个根堆

    def Insert(self, num):
        # write code here
        if not self.min_heap: # 初始比较小的min堆，这个堆是大根堆，递减
            heapq.heappush(self.min_heap, -num)
        else:
            if num > -self.min_heap[0]: # 如果数比min的最大值大，放入max堆，max堆是小根堆，递增；否则放min堆
                heapq.heappush(self.max_heap, num) 
            else:
                heapq.heappush(self.min_heap, -num)
            # 调整他俩数量差，不能大于1，一旦等于2，要平衡一下
            if len(self.min_heap) -2 == len(self.max_heap):
                cur = heapq.heappop(self.min_heap)
                heapq.heappush(self.max_heap, -cur)
            elif len(self.max_heap) -2 == len(self.min_heap):
                cur = heapq.heappop(self.max_heap)
                heapq.heappush(self.min_heap, -cur)
        
    def GetMedian(self):
        # write code here
		# 最终他俩堆的顶元素就是最中间的两个数
        if len(self.min_heap) > len(self.max_heap):
            return -self.min_heap[0]
        elif len(self.min_heap) < len(self.max_heap):
            return self.max_heap[0]
        else:
            return (-self.min_heap[0] + self.max_heap[0]) / 2