LeetCode 295. 数据流的中位数
目录
题目描述
题意分析
题目要求设计一个数据结构，支持两种操作：
addNum把一个整数加入数据流，findMedian返回当前所有已加入元素的中位数。数据是持续到来的，不是一次性给定的数组，所以不存在「先拿到全部数据再处理」的机会，每一次查询面对的都是一个随时可能变长的集合。中位数的定义要看清楚：把已加入的元素排好序后，如果个数为奇数，中位数是正中间那个数；如果个数为偶数，中位数是中间两个数的平均值。平均值意味着返回类型必须是浮点数，而不是整数。
约束里有两个明显的信号。第一，
addNum和findMedian的调用次数最多可以到 $5 \times 10^4$ 量级，两个操作交替出现，所以单次操作的代价必须足够低，不能让某一个操作退化成线性甚至更差。第二，元素值域是 $[-10^5, 10^5]$，是一个有限且不大的范围，这是一个留给进阶优化的伏笔。边界情况：数据流为空时不会调用
findMedian，所以不必处理空集合；只有一个元素时中位数就是它自己；元素可以重复，也可以是负数，中位数本身不要求是数据流中出现过的数。
解法：双堆维护中位数
核心思路
问题关键： 数据不断插入，但查询只关心排序后中间的一到两个数。每次重新排序是 $O(n \log n)$；维护有序数组虽然查询为 $O(1)$，插入时仍要移动元素，最坏为 $O(n)$。
为什么选择双堆： 把数据分成较小的一半和较大的一半，只需维护分界线附近的极值。大顶堆
small保存较小的一半，堆顶是其中最大值；小顶堆large保存较大的一半，堆顶是其中最小值。堆能在 $O(\log n)$ 内插入和移动分界元素。不变量：
small中任意元素都不大于large中任意元素；small.size == large.size，或small.size == large.size + 1。正确性： 新数不大于
small堆顶时放入small，否则放入large，有序划分仍成立。若大小失衡，只把small的最大值移到large，或把large的最小值移到small；移动的是分界元素，所以有序划分不会被破坏，同时恢复大小不变量。两堆等大时，两个堆顶正是中间两数；small多一个时，它的堆顶就是唯一中间数。
解题步骤
- 初始化大顶堆
small和小顶堆large。- 插入
num：若small为空或num <= small.peek()，放入small；否则放入large。- 调整大小：
small多两个时，把它的堆顶移到large；large更多时，把它的堆顶移到small。- 查询中位数：两堆等大就取两个堆顶的平均值，否则取
small堆顶。口述示例： 插入
5, 2, 8, 3。平衡后依次得到：small=[5]；small=[2], large=[5]；small=[5,2], large=[8]；最后small的最大值为 3、large的最小值为 5，中位数是 4。边界与反例： 只有一个数时它一定留在
small；重复值放在哪侧都可以，只要两条不变量成立。偶数个元素求平均时必须先转成浮点数，否则1, 2会错误地得到 1。
代码实现
class MedianFinder {
private final PriorityQueue<Integer> small =
new PriorityQueue<>(Collections.reverseOrder());
private final PriorityQueue<Integer> large = new PriorityQueue<>();
public MedianFinder() {
}
public void addNum(int num) {
if (small.isEmpty() || num <= small.peek()) {
small.offer(num);
} else {
large.offer(num);
}
if (small.size() > large.size() + 1) {
large.offer(small.poll());
} else if (large.size() > small.size()) {
small.offer(large.poll());
}
}
public double findMedian() {
if (small.size() == large.size()) {
return (small.peek() + (double) large.peek()) / 2;
}
return small.peek();
}
}
import "container/heap"
type IntHeap []int
func (h IntHeap) Len() int { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *IntHeap) Push(x any) {
*h = append(*h, x.(int))
}
func (h *IntHeap) Pop() any {
old := *h
last := len(old) - 1
x := old[last]
*h = old[:last]
return x
}
type MedianFinder struct {
small IntHeap // 保存相反数，用最小堆模拟最大堆
large IntHeap
}
func Constructor() MedianFinder {
return MedianFinder{}
}
func (this *MedianFinder) AddNum(num int) {
if len(this.small) == 0 || num <= -this.small[0] {
heap.Push(&this.small, -num)
} else {
heap.Push(&this.large, num)
}
if len(this.small) > len(this.large)+1 {
heap.Push(&this.large, -heap.Pop(&this.small).(int))
} else if len(this.large) > len(this.small) {
heap.Push(&this.small, -heap.Pop(&this.large).(int))
}
}
func (this *MedianFinder) FindMedian() float64 {
if len(this.small) == len(this.large) {
return (float64(-this.small[0]) + float64(this.large[0])) / 2
}
return float64(-this.small[0])
}
复杂度分析
addNum的时间复杂度为 $O(\log n)$：一次插入，至多再移动一个堆顶；findMedian只读取堆顶，为 $O(1)$。- 两个堆共同保存全部 $n$ 个元素，空间复杂度为 $O(n)$。
关键点总结
- 两个堆的堆顶要相向：较小一半取最大值，较大一半取最小值。
- “有序划分”和“大小平衡”缺一不可；只让数量接近，不能保证堆顶是中间数。
- 重平衡只能移动分界处的极值，这正是堆顶。
- 如果追问“所有数都在 0 到 100”，可以用计数数组把插入降为 $O(1)$；一般值域下双堆更通用。
易错点总结
- 只按堆大小决定新数放哪边：可能满足数量平衡，却让
small.peek() > large.peek()。- 两个堆方向写反：堆顶不再是靠近中位数的分界元素。
- 重平衡条件写成两堆必须永远等大：奇数个元素时会把唯一中间数放错位置。
- 偶数情况使用整数除法，或先用两个
int相加再转浮点：会截断小数，极端值下还可能溢出。- Go 的
small存的是相反数，比较、跨堆移动和取答案时漏掉负号都会得到错误分界。
相似题目
| 题目 | 难度 | 考察点 |
|---|---|---|
| 480. 滑动窗口中位数 | 困难 | 双堆还要支持窗口滑出时的删除，需配合延迟删除或有序集合 |
| 剑指 Offer 41. 数据流中的中位数 | 困难 | 同款双堆模板，可作为默写复现的对照题 |
| 面试题 17.20. 连续中值 | 困难 | 换了题面的同一模型，重点在快速识别「动态中位数」信号 |
| 703. 数据流中的第 K 大元素 | 简单 | 只需单个容量为 K 的小顶堆，不涉及两堆之间的平衡 |
| 215. 数组中的第K个最大元素 | 中等 | 静态数组求第 K 大，可用快速选择做到平均 $O(n)$ |