63 数据流中的中位数

题目描述

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

解题思路

核心思想：使用两个堆，最大堆存较小的一半，最小堆存较大的一半。

Python 实现

import heapq

class MedianFinder:
    def __init__(self):
        self.max_heap = []  # 存储较小的一半（用负数模拟最大堆）
        self.min_heap = []  # 存储较大的一半

    def add_num(self, num):
        """
        添加数字
        """
        # 添加到最大堆
        heapq.heappush(self.max_heap, -num)
        # 从最大堆弹出一个，添加到最小堆
        heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))

        # 如果最小堆比最大堆多1个，平衡一下
        if len(self.min_heap) > len(self.max_heap):
            heapq.heappush(self.max_heap, -heapq.heappop(self.min_heap))

    def get_median(self):
        """
        获取中位数
        """
        if not self.max_heap and not self.min_heap:
            return None

        # 奇数个，最大堆多一个，返回最大堆的堆顶
        if len(self.max_heap) > len(self.min_heap):
            return -self.max_heap[0]
        # 偶数个，返回两堆堆顶的平均
        return (-self.max_heap[0] + self.min_heap[0]) / 2


# 测试代码
if __name__ == "__main__":
    mf = MedianFinder()
    nums = [1, 2, 3, 4, 5]

    for num in nums:
        mf.add_num(num)
        median = mf.get_median()
        print(f"添加{num}, 中位数: {median}")

# 输出:
# 添加1, 中位数: 1
# 添加2, 中位数: 1.5
# 添加3, 中位数: 2
# 添加4, 中位数: 2.5
# 添加5, 中位数: 3