跳到主要内容

LeetCode 295. 数据流的中位数

本节目标

用一大一小两个堆维护数据的左右两半，在在线插入后以常数时间读取中位数。

这道题是数据结构综合应用框架中的拓展题。单个堆擅长维护一侧的极值，而中位数位于整个有序序列的中间，需要两个堆共同维护数据的左右两半。

查看原题

题意与约束​

设计一个数据结构，支持：

  • addNum(num)：向数据流加入一个整数；
  • findMedian()：返回当前所有元素的中位数。

元素个数为奇数时，中位数是排序后正中间的数；元素个数为偶数时，中位数是中间两个数的平均值。

每次查询时重新排序全部历史数据会重复大量工作。我们希望让一次插入只进行局部调整，并随时从数据结构的边界读取中位数。

两个堆分别维护左右两半​

把有序数据想象成左右两部分：

  • lower 保存较小的一半，使用大根堆，因此堆顶是左半部分最大值；
  • upper 保存较大的一半，使用小根堆，因此堆顶是右半部分最小值。

只要两个堆满足正确的不变量，中位数就在两个堆顶：

  • 总数为奇数时，让 lower 比 upper 多一个元素，中位数是 lower 堆顶；
  • 总数为偶数时，两堆大小相等，中位数是两个堆顶的平均值。

必须同时维护两个不变量​

顺序不变量​

lower 中任意元素都不大于 upper 中任意元素。

加入新数时，若它不大于 lower 堆顶，就先进入 lower；否则进入 upper。这样新元素先落在正确的一侧。

大小不变量​

始终满足：

lower.size == upper.size
或
lower.size == upper.size + 1

插入后若某一侧元素过多，就把该侧堆顶移动到另一侧。移动的是最靠近中间边界的元素，因此调整大小的同时不会破坏顺序。

只有顺序正确但数量失衡，中间位置可能不在堆顶；只有数量平衡但左右混杂，两个堆顶也不能代表中间值。这两个不变量缺一不可。

示例推演​

依次加入 5、2、8、3：

新元素lower 保存的较小一半upper 保存的较大一半中位数
5{5}{}5
2{2}{5}3.5
8{5, 2}{8}5
3{3, 2}{5, 8}4

表中只展示集合内容；真正需要立即访问的是 lower 的最大值与 upper 的最小值。

代码实现​

C++ 用一个默认大根堆和一个 greater<int> 小根堆；Python 用负数模拟大根堆。两份实现都先选择插入侧，再通过移动堆顶恢复大小不变量。

C++17
#include <functional>
#include <queue>
#include <vector>
using namespace std;

class MedianFinder {
private:
    priority_queue<int> lower;
    priority_queue<int, vector<int>, greater<int>> upper;

public:
    MedianFinder() = default;

    void addNum(int num) {
        if (lower.empty() || num <= lower.top()) {
            lower.push(num);
        } else {
            upper.push(num);
        }

        if (lower.size() > upper.size() + 1) {
            upper.push(lower.top());
            lower.pop();
        } else if (upper.size() > lower.size()) {
            lower.push(upper.top());
            upper.pop();
        }
    }

    double findMedian() {
        if (lower.size() > upper.size()) {
            return lower.top();
        }
        return lower.top() / 2.0 + upper.top() / 2.0;
    }
};

C++ 在计算偶数个元素的中位数时分别执行 lower.top() / 2.0 与 upper.top() / 2.0 后再相加，避免先做两个整数的加法可能溢出。

复杂度分析​

  • addNum 时间复杂度：O(log n)。插入一次，至多再移动一个堆顶。
  • findMedian 时间复杂度：O(1)。答案直接来自一个或两个堆顶。
  • 空间复杂度：O(n)。为支持未来更新，需要保留数据流中的全部元素，只是分布在两个堆中。

易错点​

  • 只维护两个堆大小接近，却没有保证左侧元素都不大于右侧元素。
  • 调整顺序后忘记恢复大小，导致奇数个元素时中位数不在约定的堆顶。
  • Python 忘记把大根堆中的负数还原为正数。
  • 偶数个元素时使用整数除法，丢失 .5。
  • C++ 先计算 lower.top() + upper.top()，极端整数可能在转换为浮点数之前溢出。

模式迁移​

双堆把“持续维护 Top K”扩展成“持续维护分界线两侧”。当查询目标从中位数变为某个固定百分位数时，可以调整两个堆的目标大小比例；当窗口会删除旧元素时，还需要额外记录延迟删除信息。后两类属于更复杂的维护问题，本题先掌握顺序与大小两个基本不变量即可。