LeetCode 剑指 Offer 41. 数据流中的中位数
题目描述

题意分析
要设计一个数据结构,支持两种操作:往里塞一个整数,以及随时问「目前塞进去的所有数的中位数是多少」。中位数的定义是把当前全部数字排好序后,奇数个时取正中间那个,偶数个时取中间两个的平均值,因此返回值是浮点数。
「数据流」三个字是最强的信号:数字一个个到来,总量未知,而且两种操作会任意交错出现。这排除了「攒齐再排序」的思路,也说明每次插入都必须在常数或对数时间内完成,不能等到查询时才做重活。
值得注意的是查询只关心「中间」这一两个位置,完全不关心其余元素的具体排列。这是一个很强的松弛条件——数据结构没有必要维持全序,只要能随时把中间那两个值捞出来就够了。边界上要考虑:只有一个数时中位数就是它自己;数字可以重复;数字可以为负;两个中间值相加时有超出 32 位的风险。
解法:双堆维护数据流中位数
核心思路
问题关键: 中位数只依赖有序序列中间的一到两个值,没有必要维护所有元素的完整顺序。把数据分成较小的一半和较大的一半,只需快速取得左半最大值、右半最小值。
使用最大堆
small保存较小的一半,最小堆large保存较大的一半,并维护两条不变量:
- 顺序不变量:
small中任意元素不大于large中任意元素。- 大小不变量:
small.size等于large.size,或比它多 1。插入时先按
small堆顶判断归属,再把多出边界的堆顶搬到另一侧恢复平衡。搬动的是两半的边界元素,因此不会破坏顺序。正确性: 两堆等大时,中间两个数分别是
small的最大值和large的最小值;small多一个时,它的堆顶就是唯一中间值。只要两条不变量成立,查询即可直接读堆顶。
解题步骤
- 创建最大堆
small和最小堆large。- 插入数字:
small为空或新值不大于其堆顶时放入small,否则放入large。- 若
small比large多两个,将small堆顶移到large;若large更多,将large堆顶移到small。- 查询时,奇数个元素返回
small堆顶;偶数个元素返回两个堆顶的平均值。口述示例: 依次插入 1、2 后,两堆分别为
small = {1}、large = {2},中位数为 1.5。再插入 3 时它先进入large,随后将large堆顶 2 搬到small,中位数变为 2。数值边界: Java 中两个堆顶相加前先转
long,避免整数溢出;除以2.0,避免整数除法截断。
代码实现
import java.util.Collections;
import java.util.PriorityQueue;
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 (small.size() < large.size()) {
small.offer(large.poll());
}
}
public double findMedian() {
if (small.size() > large.size()) {
return small.peek();
}
return ((long) small.peek() + large.peek()) / 2.0;
}
}
import "container/heap"
type intHeap struct {
data []int
max bool
}
func (h intHeap) Len() int { return len(h.data) }
func (h intHeap) Swap(i, j int) { h.data[i], h.data[j] = h.data[j], h.data[i] }
func (h intHeap) Less(i, j int) bool {
if h.max {
return h.data[i] > h.data[j]
}
return h.data[i] < h.data[j]
}
func (h *intHeap) Push(x any) {
h.data = append(h.data, x.(int))
}
func (h *intHeap) Pop() any {
last := len(h.data) - 1
value := h.data[last]
h.data = h.data[:last]
return value
}
type MedianFinder struct {
small intHeap
large intHeap
}
func Constructor() MedianFinder {
return MedianFinder{small: intHeap{max: true}}
}
func (m *MedianFinder) AddNum(num int) {
if m.small.Len() == 0 || num <= m.small.data[0] {
heap.Push(&m.small, num)
} else {
heap.Push(&m.large, num)
}
if m.small.Len() > m.large.Len()+1 {
heap.Push(&m.large, heap.Pop(&m.small).(int))
} else if m.small.Len() < m.large.Len() {
heap.Push(&m.small, heap.Pop(&m.large).(int))
}
}
func (m *MedianFinder) FindMedian() float64 {
if m.small.Len() > m.large.Len() {
return float64(m.small.data[0])
}
return (float64(m.small.data[0]) + float64(m.large.data[0])) / 2
}
复杂度分析
addNum时间复杂度:$O(\log n)$,包含一次入堆和至多一次跨堆搬运。findMedian时间复杂度:$O(1)$,只读取堆顶。- 空间复杂度:$O(n)$,每个已插入元素恰好位于一个堆中。
关键点总结
- 双堆只维护中位数需要的偏序,不承担完整排序的成本。
- 顺序不变量与大小不变量必须同时成立,缺一都无法正确查询。
- 约定
small可以多一个元素,可统一奇数情况。- 求平均时同时防止整数相加溢出和整数除法截断。
易错点总结
- 只按大小平衡、不保证两堆顺序:堆顶可能不是中间的两个数。
- 平衡条件写成两堆必须始终等大:奇数次插入后会把一侧搬空或反复搬运。
- 偶数情况直接做整数加法和除法:可能溢出,并丢失 0.5。
- 最大堆比较器用减法或写反:极值可能溢出,且
small堆顶不再是左半最大值。- 在题目未允许空数据查询时擅自返回 0:会掩盖调用契约错误;本题保证查询前已有数据。
相似题目
| 题目 | 难度 | 考察点 |
|---|---|---|
| 4. 寻找两个正序数组的中位数 | 困难 | 静态且已有序,用分割点二分做到对数时间,与流式场景相反 |
| 239. 滑动窗口最大值 | 困难 | 同为窗口内取极值,但单调队列比堆更契合「只求最值」的需求 |
| 295. 数据流的中位数 | 困难 | 与本题同源题面,可顺带练值域受限时的计数数组优化 |
| 480. 滑动窗口中位数 | 困难 | 在对顶堆上追加删除操作,需要延迟删除或有序多重集合 |
| 703. 数据流中的第 K 大元素 | 简单 | 同为流式分位数,但只需单个固定容量的小顶堆 |
| 面试题 17.20. 连续中值 | 困难 | 同一模型的另一份题面,适合校验平衡条件的边界写法 |