295. 数据流的中位数

力扣题目链接 (opens new window)

题目描述

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

例如 arr = [2,3,4] 的中位数是 3 。 例如 arr = [2,3] 的中位数是 (2 + 3) / 2 = 2.5 。

实现 MedianFinder 类:

• MedianFinder() 初始化 MedianFinder 对象。
• void addNum(int num) 将数据流中的整数 num 添加到数据结构中。
• double findMedian() 返回到目前为止所有元素的中位数。与实际答案相差 10^-5 以内的答案将被接受。

示例 1：

输入
["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]
输出
[null, null, null, 1.5, null, 2.0]

解释
MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1);    // arr = [1]
medianFinder.addNum(2);    // arr = [1, 2]
medianFinder.findMedian(); // 返回 1.5 ((1 + 2) / 2)
medianFinder.addNum(3);    // arr = [1, 2, 3]
medianFinder.findMedian(); // 返回 2.0

提示：

• -10^5 <= num <= 10^5
• 在调用 findMedian 之前，数据结构中至少有一个元素
• 最多 5 * 10^4 次调用 addNum 和 findMedian

思路

暴力解法是把所有数存在数组里，每次查询中位数时排序，时间复杂度O(n log n)，太慢。

核心问题：如何在动态插入数据的同时，快速获取中位数？

中位数的特点是：它把有序数组分成左右两部分，左边都 ≤ 中位数，右边都 ≥ 中位数。

那么就应该想到：用两个堆维护这两部分。

对顶堆的设计：

• 大顶堆（maxHeap）：存较小的一半元素，堆顶是这一半的最大值
• 小顶堆（minHeap）：存较大的一半元素，堆顶是这一半的最小值

维护规则：

1. 保持 maxHeap.size() == minHeap.size() 或 maxHeap.size() == minHeap.size() + 1
2. 保证 maxHeap 的所有元素 ≤ minHeap 的所有元素

这样中位数就是：

• 如果总数是奇数（maxHeap多一个），中位数是 maxHeap.top()
• 如果总数是偶数，中位数是 (maxHeap.top() + minHeap.top()) / 2

插入逻辑：

每次插入 num 时：

1. 先插入到 maxHeap
2. 把 maxHeap 的堆顶（最大值）弹出，插入到 minHeap（保证 maxHeap 的元素都 ≤ minHeap）
3. 如果此时 minHeap 的元素比 maxHeap 多，就把 minHeap 的堆顶弹出，插入到 maxHeap

这样保证了两个堆的平衡和大小关系。

模拟过程

以示例为例，依次插入 1, 2：

插入 1：

maxHeap: []
minHeap: [1]
中位数 = 1

插入 2：

maxHeap: [1]
minHeap: [2]
中位数 = (1 + 2) / 2 = 1.5

插入 3：

maxHeap: [1, (3的位置)]
minHeap: [2]

调整后：
maxHeap: [1]
minHeap: [2, 3]

再调整：
maxHeap: [1, 2]
minHeap: [3]

中位数 = 2（maxHeap多一个，取maxHeap堆顶）

实际上插入3的完整过程：

1. 3 插入 maxHeap → maxHeap=[1,3]
2. maxHeap 堆顶3弹出插入 minHeap → maxHeap=[1], minHeap=[2,3]
3. minHeap 比 maxHeap 多，minHeap 堆顶2弹出插入 maxHeap → maxHeap=[1,2], minHeap=[3]

解题代码

class MedianFinder {
private:
    priority_queue<int> maxHeap; // 大顶堆，存较小的一半
    priority_queue<int, vector<int>, greater<int>> minHeap; // 小顶堆，存较大的一半
    
public:
    MedianFinder() {}
    
    void addNum(int num) {
        maxHeap.push(num); // 先插入大顶堆
        minHeap.push(maxHeap.top()); // 把大顶堆的最大值移到小顶堆
        maxHeap.pop();
        
        // 如果小顶堆元素更多，平衡一下
        if (minHeap.size() > maxHeap.size()) {
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }
    }
    
    double findMedian() {
        if (maxHeap.size() > minHeap.size()) {
            return maxHeap.top(); // 奇数个元素
        } else {
            return (maxHeap.top() + minHeap.top()) / 2.0; // 偶数个元素
        }
    }
};

复杂度分析

• 时间复杂度：addNum 为 O(log n)，findMedian 为 O(1)
• 空间复杂度：O(n)，n 是数据流中的元素个数

其他语言

Python3

import heapq

class MedianFinder:
    def __init__(self):
        self.maxHeap = []  # 大顶堆（Python用负数模拟）
        self.minHeap = []  # 小顶堆

    def addNum(self, num: int) -> None:
        heapq.heappush(self.maxHeap, -num)  # 插入大顶堆（取负）
        heapq.heappush(self.minHeap, -heapq.heappop(self.maxHeap))  # 移到小顶堆
        
        # 平衡两个堆
        if len(self.minHeap) > len(self.maxHeap):
            heapq.heappush(self.maxHeap, -heapq.heappop(self.minHeap))

    def findMedian(self) -> float:
        if len(self.maxHeap) > len(self.minHeap):
            return -self.maxHeap[0]
        else:
            return (-self.maxHeap[0] + self.minHeap[0]) / 2.0

Java

class MedianFinder {
    private PriorityQueue<Integer> maxHeap; // 大顶堆
    private PriorityQueue<Integer> minHeap; // 小顶堆

    public MedianFinder() {
        maxHeap = new PriorityQueue<>((a, b) -> b - a);
        minHeap = new PriorityQueue<>();
    }
    
    public void addNum(int num) {
        maxHeap.offer(num);
        minHeap.offer(maxHeap.poll());
        
        if (minHeap.size() > maxHeap.size()) {
            maxHeap.offer(minHeap.poll());
        }
    }
    
    public double findMedian() {
        if (maxHeap.size() > minHeap.size()) {
            return maxHeap.peek();
        } else {
            return (maxHeap.peek() + minHeap.peek()) / 2.0;
        }
    }
}

Go

import "container/heap"

type MaxHeap []int
func (h MaxHeap) Len() int           { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h MaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MaxHeap) Push(x interface{}) { *h = append(*h, x.(int)) }
func (h *MaxHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

type MinHeap []int
func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x interface{}) { *h = append(*h, x.(int)) }
func (h *MinHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

type MedianFinder struct {
    maxHeap *MaxHeap
    minHeap *MinHeap
}

func Constructor() MedianFinder {
    maxH := &MaxHeap{}
    minH := &MinHeap{}
    heap.Init(maxH)
    heap.Init(minH)
    return MedianFinder{maxH, minH}
}

func (this *MedianFinder) AddNum(num int) {
    heap.Push(this.maxHeap, num)
    heap.Push(this.minHeap, heap.Pop(this.maxHeap))
    
    if this.minHeap.Len() > this.maxHeap.Len() {
        heap.Push(this.maxHeap, heap.Pop(this.minHeap))
    }
}

func (this *MedianFinder) FindMedian() float64 {
    if this.maxHeap.Len() > this.minHeap.Len() {
        return float64((*this.maxHeap)[0])
    }
    return float64((*this.maxHeap)[0]+(*this.minHeap)[0]) / 2.0
}

JavaScript

var MedianFinder = function() {
    this.maxHeap = new MaxPriorityQueue(); // 大顶堆
    this.minHeap = new MinPriorityQueue(); // 小顶堆
};

MedianFinder.prototype.addNum = function(num) {
    this.maxHeap.enqueue(num);
    this.minHeap.enqueue(this.maxHeap.dequeue().element);
    
    if (this.minHeap.size() > this.maxHeap.size()) {
        this.maxHeap.enqueue(this.minHeap.dequeue().element);
    }
};

MedianFinder.prototype.findMedian = function() {
    if (this.maxHeap.size() > this.minHeap.size()) {
        return this.maxHeap.front().element;
    } else {
        return (this.maxHeap.front().element + this.minHeap.front().element) / 2.0;
    }
};

与代码随想录联系

本题是堆的经典应用。对顶堆是一种常见的技巧，适用于动态维护中位数、第k大元素等场景。

在代码随想录堆章节中，讲解了：

• 215. 数组中的第K个最大元素
• 347. 前 K 个高频元素

堆的核心：快速找最值。

对顶堆把数据分成两部分，用两个堆分别维护，是"分治"思想的体现。

如果录友们对堆还不熟悉，建议先学习单个堆的应用（比如第K大），再来理解对顶堆会更容易。