堆排序 之实现数据流中的中位数

实现语言：Python3.9

题目来源：牛客

实现步骤：

• 1、使用堆的方式实现，具体实现思路：我们把数据从中点位置分为两个部分，前一部分构建成大顶堆A，后一部分构建成小顶堆B（注意前半部分的数据小于后半部分的数据）；
• 2、先将数据添加到大顶堆中，然后取出大顶堆的堆顶元素，把堆顶元素添加到小顶堆中，这样做的目的是保证大顶堆中的数据全部小于小顶堆中的数据；
• 3、然后平衡A和B的数据，如果A中的数据个数小于B那么就把B的堆顶元素取出放入到A中，总之就是需要保证len(A)>=len(B)；
• 4、这样如果len(A) > len(B)，那么中点就是A部分的堆顶数据，如果len(A) == len(B) 那么中点就是A和B的堆顶数据和的一半。

import heapq
class Solution:
    min_heap = []
    max_heap = []

    # 构建大顶堆
    def my_max_heapify(self, it):
        max_heap_ = [(-x, x) for x in it]
        heapq.heapify(max_heap_)
        return [x for (_, x) in max_heap_]

    def Insert(self, num):
        # write code here
        # 先加入大顶堆中
        self.max_heap.append(num)
        # 构建大顶堆
        self.max_heap = self.my_max_heapify(self.max_heap)
        # 将大顶堆中的堆顶元素取出，放到小顶堆中
        heapq.heappush(self.min_heap, heapq.heappop(self.max_heap))
        # 再次构建大顶堆
        self.max_heap = self.my_max_heapify(self.max_heap)

        # 平衡两个堆中的元素数量
        # 如果大顶堆中的元素比小顶堆中的元素少，则将小顶堆的堆顶元素取出放到大顶堆。
        if len(self.max_heap) < len(self.min_heap):
            self.max_heap.append(heapq.heappop(self.min_heap))
            self.max_heap = self.my_max_heapify(self.max_heap)
        
    def GetMedian(self):
        # write code here
        # 如果大顶堆的元素比小顶堆的元素个数多，说明大顶堆的堆顶元素就是中间值
        if len(self.max_heap) > len(self.min_heap):
            middle = self.max_heap[0]
            return middle
        # 如果相等，则中间元素是大顶堆和小顶堆的堆顶元素
        else:
            a =  self.max_heap[0]
            b =  self.min_heap[0]
            return (a + b) / 2