目录

题目描述

剑指 Offer 41. 数据流中的中位数

image-20241107211154254

题意分析

要设计一个数据结构,支持两种操作:往里塞一个整数,以及随时问「目前塞进去的所有数的中位数是多少」。中位数的定义是把当前全部数字排好序后,奇数个时取正中间那个,偶数个时取中间两个的平均值,因此返回值是浮点数。

「数据流」三个字是最强的信号:数字一个个到来,总量未知,而且两种操作会任意交错出现。这排除了「攒齐再排序」的思路,也说明每次插入都必须在常数或对数时间内完成,不能等到查询时才做重活。

值得注意的是查询只关心「中间」这一两个位置,完全不关心其余元素的具体排列。这是一个很强的松弛条件——数据结构没有必要维持全序,只要能随时把中间那两个值捞出来就够了。边界上要考虑:只有一个数时中位数就是它自己;数字可以重复;数字可以为负;两个中间值相加时有超出 32 位的风险。

解法:双堆维护数据流中位数

核心思路

问题关键: 中位数只依赖有序序列中间的一到两个值,没有必要维护所有元素的完整顺序。把数据分成较小的一半和较大的一半,只需快速取得左半最大值、右半最小值。

使用最大堆 small 保存较小的一半,最小堆 large 保存较大的一半,并维护两条不变量:

  1. 顺序不变量: small 中任意元素不大于 large 中任意元素。
  2. 大小不变量: small.size 等于 large.size,或比它多 1。

插入时先按 small 堆顶判断归属,再把多出边界的堆顶搬到另一侧恢复平衡。搬动的是两半的边界元素,因此不会破坏顺序。

正确性: 两堆等大时,中间两个数分别是 small 的最大值和 large 的最小值;small 多一个时,它的堆顶就是唯一中间值。只要两条不变量成立,查询即可直接读堆顶。

解题步骤

  1. 创建最大堆 small 和最小堆 large
  2. 插入数字:small 为空或新值不大于其堆顶时放入 small,否则放入 large
  3. smalllarge 多两个,将 small 堆顶移到 large;若 large 更多,将 large 堆顶移到 small
  4. 查询时,奇数个元素返回 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. 连续中值 困难 同一模型的另一份题面,适合校验平衡条件的边界写法