LeetCode 剑指 Offer 41. 数据流中的中位数
题目描述
题意分析
数据流会不断加入整数，需要在任意一次加入后查询当前全部数据的中位数。中位数按数值排序后的排名定义，与数字到达的先后顺序无关，也不是所有数的平均值。
当前数量为奇数时，取排序后正中间的一个数；为偶数时，取中间两个数的平均，因此查询返回浮点数。相同数字的每次加入都独立计数，查询不会删除数据。查询中位数时应已经至少加入一个元素。
解法：双堆维护数据流中位数
核心思路
[!blue]
中位数只依赖排序后两半数据的交界，不需要每次都把所有数据完全排序。用大根堆
small保存较小的一半，堆顶是这一半的最大值；用小根堆large保存较大的一半，堆顶是这一半的最小值。只维护边界所需的堆序，就能直接取得中间值。需要同时保持两条不变量：
small中所有数不大于large中所有数；small的元素数等于large，或者恰好多一个。第一条保证两堆对应排序后的左右两半，第二条保证它们的交界正好位于整体中间。加入一个新数时，若
small为空或新数不大于它的堆顶，就放入small；否则放入large。前一种情况的新数不会大于右半任何数，后一种情况的新数不会小于左半任何数，因此值域顺序仍然成立，只可能破坏数量平衡。若左半多出两个，就将左半最大值移到右半；若右半比左半更多，就将右半最小值移到左半。移动的是两半的边界值，剩余左半仍不大于右半，数量也恢复为相等或左多一。每次只加入一个数，原来的数量差至多为一，所以最多移动一个堆顶就足够。
查询时，左半多一说明总数为奇数，它的最大值就是正中间的元素；两堆等大说明总数为偶数，中间两个值分别是左半最大值与右半最小值，取平均即可。求平均前先扩大数值类型，避免整数加法溢出或除法丢掉小数。
解题步骤
- 创建保存较小一半的大根堆
small，以及保存较大一半的小根堆large。- 加入新数时，与
small的堆顶比较，决定先放入哪一半；左堆为空时直接放入左堆。- 若
small.size > large.size + 1，将左堆最大值移到右堆；若small.size < large.size，将右堆最小值移到左堆。- 查询时，左堆多一个元素就返回它的堆顶；否则返回两个堆顶的平均值。
- 数据持续保留在两堆中，后续加入与查询继续使用相同的不变量。
代码实现
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
}
复杂度分析
- 时间复杂度：添加一个数为均摊
O(log(n + 1))，其中n是添加前的数据量。一次添加及至多一次跨堆移动只涉及常数次堆操作；查询只读取堆顶，为O(1)。- 空间复杂度：
O(n)。两堆合计保存全部已加入的元素，每次出现只存放在一侧。
关键点总结
[!green]
- 两堆只维护排名分界，不要求堆内部形成完整排序。
- 值域有序与数量平衡必须同时成立，缺少任意一条都无法保证堆顶是中位数。
- 调整大小时只移动边界堆顶，既恢复数量，也保持左右值域顺序。
- 奇数时固定让左半多一个，使查询规则保持一致。
易错点总结
[!yellow]
- 左右堆类型颠倒：较小一半需要立即取最大值，较大一半需要立即取最小值。
- 只按数量分配新数：两堆大小正确不代表左边所有值都小于右边，分侧时还需按数值边界判断。
- 要求两堆永远一样大：总数为奇数时必须有一边多一个，当前实现约定是左边。
- 重新平衡时随意搬一个元素：应移动左堆最大值或右堆最小值，否则可能破坏值域顺序。
- 先做整数加法再转换类型：溢出可能在转换前已经发生，应先提升至少一个操作数；平均值也不能用整数除法截断。
- 把查询实现成弹堆：查询只读取中位数，不应移除数据或破坏两堆平衡。
相似题目
| 题目 | 难度 | 关联与区别 |
|---|---|---|
| 480. 滑动窗口中位数 | 困难 | 中位数维护相同，滑动窗口还需删除过期值或做延迟删除。 |
| 703. 数据流中的第 K 大元素 | 简单 | 同样维护排名边界，原题k固定可只保留k个最大值，本题中间位置随总数变化。 |
转载与许可
许可
本博客所有文章除特别声明外，均采用 CC BY-NC-SA 4.0 许可协议，转载请注明出处！