动态中位数

动态中位数

问题的本质

不断往一个集合里追加数字，每加一个数之后，立刻知道当前所有数的中位数。

朴素做法：每次追加后排序，取中间。$n$ 次操作，每次 $O(n \log n)$，总共 $O(n^2 \log n)$。

瓶颈很明显——每次都全量排序，但前 $n-1$ 个数的顺序信息在上一步已经知道了，完全浪费了。

关键洞察：求中位数不需要知道整个数组的完整顺序，只需要把数分成较小的一半和较大的一半，然后看两半的交界。

数据结构：两个堆

维护两个堆：

堆	类型	存什么	堆顶是什么
lo	大根堆	较小的一半	较小一半中的最大值
hi	小根堆	较大的一半	较大一半中的最小值

始终维持一个不变量：lo 的大小等于 hi 的大小，或多 1。

在这个不变量下：

• 元素总数为奇数 → 中位数就是 lo 堆顶（较小一半的最大值）
• 元素总数为偶数 → 中位数是 lo 堆顶和 hi 堆顶的平均

$$\text{median} = \begin{cases} \text{lo.top()}, & |\text{lo}| > |\text{hi}| \\ \frac{\text{lo.top()} + \text{hi.top()}}{2}, & |\text{lo}| = |\text{hi}| \end{cases}$$

插入操作

新来一个数 x，怎么放进去又不破坏不变量？

分两步走：

1. 决定去哪个堆：x ≤ lo 堆顶就放进 lo（它属于较小的一半），否则放进 hi
2. 平衡两个堆的大小：如果 lo 比 hi 多超过 1 个，就把 lo 堆顶搬到 hi；反过来 hi 比 lo 多，就把 hi 堆顶搬到 lo

void add(int x) {
    if (lo.empty() || x <= lo.top())
        lo.push(x);
    else
        hi.push(x);

    // 平衡：lo 最多比 hi 多 1
    if (lo.size() > hi.size() + 1) {
        hi.push(lo.top());
        lo.pop();
    } else if (hi.size() > lo.size()) {
        lo.push(hi.top());
        hi.pop();
    }
}

每次插入最多触发两次堆操作（一次 push + 可能的一次搬迁），每次 $O(\log n)$。

完整实现

C++ 中 priority_queue 默认是大根堆，小根堆取反即可。

#include <queue>
#include <vector>
using namespace std;

class MedianFinder {
    priority_queue<int> lo;                          // 大根堆
    priority_queue<int, vector<int>, greater<int>> hi;  // 小根堆

public:
    void addNum(int x) {
        if (lo.empty() || x <= lo.top())
            lo.push(x);
        else
            hi.push(x);

        if (lo.size() > hi.size() + 1) {
            hi.push(lo.top()); lo.pop();
        } else if (hi.size() > lo.size()) {
            lo.push(hi.top()); hi.pop();
        }
    }

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

为什么这样是对的

核心是那个不变量：lo 存了较小的一半，hi 存了较大的一半，lo 最多多一个。

插入时不变量为什么不会断？

• 第一步把 x 放到正确的半边（根据和 lo.top() 的比较），所以 lo 里所有数 ≤ hi 里所有数，这个有序关系始终成立
• 第二步搬迁堆顶，搬过去的一定是原来那一半的边界元素，所以有序关系仍然成立
• 搬迁之后两堆大小差 ≤ 1，不变量恢复

复杂度

操作	时间
插入	$O(\log n)$
查询中位数	$O(1)$
$n$ 次追加 + 查询	$O(n \log n)$

对比朴素做法的 $O(n^2 \log n)$，双堆法把问题从”每次全量排序”降到了”每次对数级别的调整”。

示例

依次插入 $5, 1, 3, 2, 4$：

步骤	插入	lo（大根堆）	hi（小根堆）	中位数
1	5	{5}	{}	5
2	1	{1}	{5}	$(1+5)/2 = 3$
3	3	{3, 1}	{5}	3
4	2	{2, 1}	{3, 5}	$(2+3)/2 = 2.5$
5	4	{3, 2, 1}	{4, 5}	3

注意第 4 步：2 放进 lo 后 lo 有 3 个、hi 有 1 个，差了 2，所以把 lo 堆顶 3 搬到 hi。