CS-Notes 剑指 Offer 41.1 数据流中的中位数：用两个堆在线维护中位数

本文基于 CS-Notes 仓库中的题解文档 41.1 数据流中的中位数，讲解如何在元素逐个到达、无法一次性排序的数据流场景下，用"大顶堆 + 小顶堆"的对偶结构把中位数查询做到 O(1)、插入做到 O(log N)。读完本篇，你将掌握该题完整的 Java 实现、两个堆必须维持的不变式证明思路，以及它与仓库中同系列的最小的 K 个数共用的大顶堆技巧。

题目描述

本题出自《剑指 Offer》第 41 题的第一问，属于仓库剑指 Offer 题解 - 目录中"栈队列堆"分类下的典型题目（原题在线练习平台为牛客网）。

如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。

题目给出了两个操作接口：

• Insert(Integer val)：数据流来一个新元素，需要插入；
• GetMedian()：随时查询当前已读入全部元素的中位数。

难点在于**"流式 + 任意时刻可查询"**：元素是陆续到达的，任何时刻都可能被要求给出中位数，因此不能"来一批再排序"，必须用某种数据结构在插入时就维护好中位数所需的信息。

解题思路：大顶堆存左半边，小顶堆存右半边

朴素做法是维护一个有序数组，每次插入 O(N)、查询 O(1)；或者每次查询时现排，插入 O(1)、查询 O(N log N)。两种方案都偏科，无法同时满足"插入快、查询也快"。

官方题解采用的核心思想是：用两个堆把数据流"对半切"，让中位数永远落在两个堆顶的交汇处：

成员变量	类型	职责
left	PriorityQueue<Integer>，比较器 (o1, o2) -> o2 - o1，即大顶堆	存储排序后左半边（较小的一半）元素，堆顶是左半边的最大值
right	PriorityQueue<Integer>，默认比较器，即小顶堆	存储排序后右半边（较大的一半）元素，堆顶是右半边的最小值
N	int	当前数据流已读入的元素个数

这里涉及的两个堆，正是仓库 Java 容器 中介绍的 Queue 分类下的 PriorityQueue——"基于堆结构实现，可以用它来实现优先队列"，默认是小顶堆，初始化时传入比较器即可反转为大顶堆。

两个必须维持的不变式

整个算法正确性依赖以下两条不变式（invariant）：

1. 有序性不变式：right 中的每个元素都大于等于 left 中的每个元素。这样两个堆顶就是整个数据集排序后"正中间"的 1 个或 2 个元素；
2. 规模平衡不变式：两堆大小相差不超过 1，且奇偶关系由 N 决定——N 为奇数时 right 比 left 多一个元素，N 为偶数时两堆等大。

维持住这两条不变式后，中位数查询就是 O(1)：

• N 为奇数：中位数是排序后的中间一个数，恰好是 right 的堆顶（右半边的最小值）；
• N 为偶数：中位数是中间两数的平均值，恰好是 left.peek() + right.peek() 的平均值。

Insert：先插入"错误"的一边，再把堆顶挪到另一边

Insert 的写法看似绕，其实是在利用堆顶特性以 O(log N) 完成"跨堆搬运"，从而保证有序性不变式：

/* 大顶堆，存储左半边元素 */
private PriorityQueue<Integer> left = new PriorityQueue<>((o1, o2) -> o2 - o1);
/* 小顶堆，存储右半边元素，并且右半边元素都大于左半边 */
private PriorityQueue<Integer> right = new PriorityQueue<>();
/* 当前数据流读入的元素个数 */
private int N = 0;

public void Insert(Integer val) {
    /* 插入要保证两个堆存于平衡状态 */
    if (N % 2 == 0) {
        /* N 为偶数的情况下插入到右半边。
         * 因为右半边元素都要大于左半边，但是新插入的元素不一定比左半边元素来的大，
         * 因此需要先将元素插入左半边，然后利用左半边为大顶堆的特点，
         * 取出堆顶元素即为最大元素，此时插入右半边 */
        left.add(val);
        right.add(left.poll());
    } else {
        right.add(val);
        left.add(right.poll());
    }
    N++;
}

逐分支解释（这是原文档给出的完整实现，逐行注释如下）：

• N 为偶数（此时两堆等大，插入后应由 right 变多一个）：目标边是 right，但新元素 val 完全可能比 left 里某些元素还小，不能直接扔进 right。于是先把 val 塞进 left（大顶堆），此时 left 的堆顶就是"左半边 + val"中的最大值；把它 poll 出来放进 right。这个最大值正是排序后应该属于右半边的最小那个数，一次 add + 一次 poll 就同时完成了"入堆"与"修正归属"；
• N 为奇数（此时 right 多一个元素，插入后应由 left 变多一个）：对称操作——先 right.add(val)，再把 right 堆顶（右半边最小值）搬进 left。

每轮插入都恰好发生一次"跨堆搬运"，所以规模平衡不变式天然成立，无需额外的 rebalance 逻辑。

GetMedian：只读堆顶，O(1) 返回

public Double GetMedian() {
    if (N % 2 == 0)
        return (left.peek() + right.peek()) / 2.0;
    else
        return (double) right.peek();
}

• N 为偶数：两堆等大，中间两个数就是 left 的最大值与 right 的最小值。注意除以 2.0 触发浮点除法，避免整数截断；
• N 为奇数：中间那个数就是 right 的堆顶，强转 double 后返回。

两个方法都只用到 peek（O(1)），配合 Insert 的两次堆操作，得到整体复杂度：

操作	时间复杂度	说明
Insert	O(log N)	一次 add + 一次 poll，均为 O(log N)
GetMedian	O(1)	仅读堆顶
空间	O(N)	N 个元素全部分布在两个堆中

对比"数组排序"方案（查询 O(1) 但插入 O(N)，或查询 O(N)），双堆方案把两个操作都压到对数以内，这正是它适合数据流场景的原因。

执行过程走查：以插入序列 5、2、3、4 为例

按上面的代码逐步跟踪，可以直观看到不变式如何被维持：

步骤	操作	left（大顶堆，堆顶在前）	right（小顶堆，堆顶在前）	N 后	GetMedian	验证（全量排序）
1	Insert(5)	[5] → 搬出 5	[5]	1	5.0	[5] → 5 ✓
2	Insert(2)	[2]（2 先入 left，再与 5 比较后留下）	[5]	2	(2+5)/2 = 3.5	[2,5] → 3.5 ✓
3	Insert(3)	[3, 2]（3 先入 left，3 出堆顶）	[3, 5]	3	3.0	[2,3,5] → 3 ✓
4	Insert(4)	[3, 2]（4 先入 right，3 出堆顶）	[4, 5]	4	(3+4)/2 = 3.5	[2,3,4,5] → 3.5 ✓

注意第 2 步：val = 2 小于左半边已有的 5，它先入 left 后堆顶仍是 5，5 被搬到 right，2 留在 left——这正是"新元素不一定比左半边大"时该分支存在的意义。第 4 步则相反，val = 4 进入 right 后堆顶是最小的 3，3 被搬回 left。无论哪种情况，搬动的总是"当前跨界的那个元素"。

细节与易错点

1. 比较器溢出隐患：仓库题解（与 40. 最小的 K 个数 一致）用 (o1, o2) -> o2 - o1 实现大顶堆。做减法在极端值下可能溢出（例如 o1 = Integer.MIN_VALUE），更稳妥的等价写法是 (o1, o2) -> Integer.compare(o2, o1)，面试中说明这一点是加分项。
2. 偶数个数时除的是 2.0：(left.peek() + right.peek()) / 2.0，若写成 / 2 会做整数除法截断小数。
3. 堆的归属只约束"跨界"元素：本方案从不移动堆内普通元素，每次插入只做一次堆顶搬运，这是 O(log N) 复杂度的来源，也是它与"两个有序列表手动 rebalance"方案的本质区别。
4. 空流与查询时机：接口约定 GetMedian 在已插入元素后调用，N 计数即用来判断当前属于奇数还是偶数情形，无需额外调用 size()。

为什么不用别的方案

• 数组 + 每次排序：查询快但插入 O(N)，数据流持续增长时总代价劣于堆；
• 平衡二叉搜索树（如 TreeSet）：Java 的 TreeSet（基于红黑树，见 Java 容器 中 Set 分类）也能做到插入/查询 O(log N)，但要取"第 N/2 个"位置仍需额外维护索引或顺序统计；双堆方案更简单，且只关心中间一个/两个位置，堆顶恰好就是答案，是最贴合题意的结构；
• 快速选择（QuickSelect）：单次求中位数 O(N)，但无法在每次插入后低成本复用，不适合"任意时刻查询"的流式场景。仓库中快速选择的思想可见 40. 最小的 K 个数 的第二个方案，适合"一次性离线求解"，与本在线场景正好互补。

仓库中的相关题解

• 40. 最小的 K 个数：与本题共用"大顶堆 + 堆顶即当前极值"的技巧，该题解还专门强调"应该使用大顶堆来维护最小堆"以及 PriorityQueue 比较器的用法；
• 41.2 字符流中第一个不重复的字符：同属"流式插入 + 即时查询"题型，用频次数组 + 队列维护答案，可与双堆方案对照体会"流式问题"的通用解法思路；
• 剑指 Offer 题解 - 目录：本题归类于"栈队列堆"章节，可顺藤摸瓜练习同分类的9. 用两个栈实现队列、59. 滑动窗口的最大值等。

小结

数据流中位数问题的关键是把"中位数"转化为"两个堆顶"：left 大顶堆存较小一半、right 小顶堆存较大一半，插入时通过"先入目标侧的对面、再搬堆顶"的一次跨堆操作同时完成入堆与归属修正，从而在 O(log N) 插入、O(1) 查询下任意时刻返回正确中位数。完整可运行代码见仓库题解 41.1 数据流中的中位数，配合上文走查表验证每一步后，可以直接在刷题平台上实现 Insert 与 GetMedian 两个方法。