LeetCode 295. Find Median from Data Stream 题解

LeetCode 295. Find Median from Data Stream 题解

题目描述

中位数是有序整数列表中的中间值。如果列表的大小是偶数，则没有中间值，中位数是两个中间值的平均值。

例如，[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]

解题思路

使用两个堆来维护数据：

• 大顶堆（max_heap）：存储较小的一半数据
• 小顶堆（min_heap）：存储较大的一半数据

保持两个堆的大小差不超过1，这样中位数就在堆顶。

代码实现

import heapq

class MedianFinder:
    def __init__(self):
        # 小顶堆，存储较大的一半
        self.min_heap = []
        # 大顶堆，存储较小的一半（Python用负数模拟）
        self.max_heap = []
    
    def addNum(self, num):
        # 先加入大顶堆
        heapq.heappush(self.max_heap, -num)
        
        # 将大顶堆的最大值移到小顶堆
        max_val = -heapq.heappop(self.max_heap)
        heapq.heappush(self.min_heap, max_val)
        
        # 平衡两个堆的大小
        if len(self.min_heap) > len(self.max_heap):
            min_val = heapq.heappop(self.min_heap)
            heapq.heappush(self.max_heap, -min_val)
    
    def findMedian(self):
        if len(self.max_heap) > len(self.min_heap):
            return -self.max_heap[0]
        else:
            return (-self.max_heap[0] + self.min_heap[0]) / 2.0

复杂度分析

• 时间复杂度：
  • addNum：O(log n)
  • findMedian：O(1)
• 空间复杂度：O(n)

总结

本题是双堆技巧的经典应用。

关键点：

• 使用两个堆维护数据的两半
• 保持堆的大小平衡
• 中位数在堆顶