LeetCode 295. 数据流的中位数——双堆维护中位数

题目描述

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

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

中位数的定义是：

• 如果元素个数是奇数，中位数是排序后最中间的数。
• 如果元素个数是偶数，中位数是排序后中间两个数的平均值。

例如：

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

这道题的难点在于：数据是不断加入的，不能每次 findMedian() 都重新排序。

核心思路

用两个堆把所有数分成左右两半：

l：大顶堆，保存较小的一半
r：小顶堆，保存较大的一半

其中：

l.peek() = 左半部分最大值
r.peek() = 右半部分最小值

这样中位数就只和两个堆顶有关：

如果总数是奇数：中位数 = l.peek()
如果总数是偶数：中位数 = (l.peek() + r.peek()) / 2.0

为什么这样分

如果把所有数字排好序，中位数只会出现在中间位置。

所以不需要维护完整有序数组，只需要维护：

较小的一半 | 较大的一半

左边最大的数和右边最小的数，正好就是中间附近的两个数。

大顶堆适合快速拿到左半部分最大值；小顶堆适合快速拿到右半部分最小值。

关键不变量

整个过程中要始终维护两个条件：

1. l 中的数整体 <= r 中的数
2. l.size() == r.size()，或者 l.size() == r.size() + 1

也就是说，l 要么和 r 一样多，要么只比 r 多一个。

为什么让 l 多一个？

因为这样元素个数为奇数时，可以直接返回：

l.peek()

不用再判断中位数在哪个堆里。

addNum 的过程

代码中有两个分支。

情况一：两个堆大小相同

if (l.size() == r.size()) {
    r.offer(num);
    l.offer(r.poll());
}

此时加入新元素后，应该让 l 比 r 多一个。

操作过程是：

1. 先把 num 放进 r
2. 再把 r 中最小的数移动到 l

这样可以保证移动到 l 的数，仍然属于较小的一半。

情况二：l 比 r 多一个

else {
    l.offer(num);
    r.offer(l.poll());
}

此时加入新元素后，应该让两个堆重新变成一样大。

操作过程是：

1. 先把 num 放进 l
2. 再把 l 中最大的数移动到 r

这样可以保证移动到 r 的数，属于较大的一半。

手动模拟

以依次加入 1, 2, 3, 4, 5 为例。

这里展示的是逻辑顺序，不代表 Java PriorityQueue 的真实内部数组顺序。

初始：

l = []
r = []

addNum(1)

两个堆大小相同，先放入 r，再把 r 的最小值移动到 l：

l = [1]
r = []

当前中位数：

1

addNum(2)

l 比 r 多一个，先放入 l，再把 l 的最大值移动到 r：

l = [1]
r = [2]

当前中位数：

(1 + 2) / 2.0 = 1.5

addNum(3)

两个堆大小相同：

先 r.offer(3)，r = [2,3]
再 l.offer(r.poll())，把 2 放入 l

结果：

l = [2,1]
r = [3]

当前中位数：

2

addNum(4)

l 比 r 多一个：

先 l.offer(4)，l = [4,1,2]
再 r.offer(l.poll())，把 4 放入 r

结果：

l = [2,1]
r = [3,4]

当前中位数：

(2 + 3) / 2.0 = 2.5

addNum(5)

两个堆大小相同：

先 r.offer(5)，r = [3,4,5]
再 l.offer(r.poll())，把 3 放入 l

结果：

l = [3,1,2]
r = [4,5]

当前中位数：

3

所以中位数依次是：

1, 1.5, 2, 2.5, 3

Java 代码

class MedianFinder {
    PriorityQueue<Integer> l;
    PriorityQueue<Integer> r;

    public MedianFinder() {
        l = new PriorityQueue<>((a, b) -> b - a);
        r = new PriorityQueue<>();
    }

    public void addNum(int num) {
        if (l.size() == r.size()) {
            r.offer(num);
            l.offer(r.poll());
        } else {
            l.offer(num);
            r.offer(l.poll());
        }
    }

    public double findMedian() {
        if (l.size() > r.size()) return l.peek();
        return (l.peek() + r.peek()) / 2.0;
    }
}

PriorityQueue 操作区分

这题里最重要的是理解 peek() 和 poll()：

peek()：查看堆顶元素，不删除
poll()：取出堆顶元素，并删除
offer()：加入一个元素

默认的 PriorityQueue 是小顶堆：

r = new PriorityQueue<>();

所以：

r.peek() 是 r 中的最小值

如果传入比较器：

l = new PriorityQueue<>((a, b) -> b - a);

就可以让 l 变成大顶堆：

l.peek() 是 l 中的最大值

易错点

1. PriorityQueue 的队头不是最早加入的元素，而是优先级最高的元素。
2. 默认 PriorityQueue 是小顶堆，不是大顶堆。
3. peek() 只是查看堆顶，poll() 才会删除堆顶。
4. 只维护数量平衡还不够，还要保证 l.peek() <= r.peek()。
5. 偶数个数时要写 / 2.0，否则容易写成整数除法。
6. 当前写法让 l 始终不少于 r，所以奇数时返回的是 l.peek()。

复杂度分析

每次加入一个数时，会进行堆的插入和删除：

addNum：O(log n)

查找中位数只需要看堆顶：

findMedian：O(1)

两个堆一共保存所有元素：

空间复杂度：O(n)

总结

这道题的核心不是排序，而是维护中位数两边的边界。

用大顶堆 l 保存较小的一半，用小顶堆 r 保存较大的一半。只要维护好两个不变量：

l 中的数整体 <= r 中的数
l.size() == r.size() 或 l.size() == r.size() + 1

中位数就可以通过堆顶快速得到。

下次重写前可以先问自己：

1. 两个堆分别保存哪一半？
2. l.peek() 和 r.peek() 分别代表什么？
3. 为什么加入元素时要先放进一个堆，再把堆顶移动到另一个堆？
4. 奇数个元素时，为什么可以直接返回 l.peek()？