295 数据流的中位数

题目描述#

咱来翻译下题面：需要实现一个 MedianFinder，它会不断接收数据流里的整数，并且随时返回当前所有数字的中位数。

如果当前数字个数是奇数，中位数就是排序后正中间的数字；如果是偶数，中位数就是中间两个数字的平均值。也就是说，addNum 负责把新数加进来，findMedian 要在已经加入的所有数里快速找到中位数。

思路#

这题要是每次 findMedian 都排序，哎呀，数据流一长就顶不住啦。中位数真正关心的是“较小的一半里最大的数”和“较大的一半里最小的数”，所以咱可以用两个堆把这两边分开维护。

左边的堆保存较小的一半，但 Go 标准库只有小根堆，所以代码里把左边元素取相反数入堆，这样 left.IntSlice[0] 对应的负数越小，原值就越大，相当于模拟了一个大根堆。右边的堆直接保存较大的一半，是正常小根堆。这样左堆堆顶能拿到较小一半的最大值，右堆堆顶能拿到较大一半的最小值。

插入时要同时维护两个不变式：第一，左堆元素数量要么等于右堆，要么比右堆多一个；第二，左边所有原始值都不大于右边所有值。代码的做法很巧：如果两个堆一样大，就先把新数放进右堆，再把右堆最小值移到左堆；如果左堆更大，就先把新数放进左堆，再把左堆最大值移到右堆。这样每次插入后，两边数量和大小关系都会自动调好。

查询中位数就简单啦：如果左堆更多，答案就是左堆堆顶还原后的值；如果两边一样多，答案就是左堆最大值和右堆最小值的平均数。每次插入需要 O(log n)，查询只要 O(1)，空间复杂度是 O(n)。

知识边界#

双堆维护中位数的核心是不变式：左半部分最大值不超过右半部分最小值，并且两边长度差最多为 1。
Go 的 container/heap 默认依赖 sort.IntSlice 的小根堆行为；想模拟大根堆，可以把值取相反数存进去，取出时再还原。
平衡两个堆时，不能只看长度，还要把跨过分界线的堆顶元素转移过去，否则左右两边的大小关系可能被新插入的数破坏。
偶数个元素时，中位数是两个堆顶的平均值；本题代码里左堆存的是负数，所以表达式 rightTop - leftTopNeg 实际上就是 rightTop + leftTopOriginal。

代码#

代码就按“双堆夹住中位数”的思路来，左边负责守住小半边的最大值，右边负责守住大半边的最小值——

import (
 "container/heap"
 "sort"
)

type MedianFinder struct {
 left  hp // 入堆的元素取相反数，变成最大堆
 right hp // 最小堆
}

func Constructor() (_ MedianFinder) {
 return
}

func (mf *MedianFinder) AddNum(num int) {
 if mf.left.Len() == mf.right.Len() {
  heap.Push(&mf.right, num)
  heap.Push(&mf.left, -heap.Pop(&mf.right).(int))
 } else {
  heap.Push(&mf.left, -num)
  heap.Push(&mf.right, -heap.Pop(&mf.left).(int))
 }
}

func (mf *MedianFinder) FindMedian() float64 {
 if mf.left.Len() > mf.right.Len() {
  return float64(-mf.left.IntSlice[0])
 }
 return float64(mf.right.IntSlice[0]-mf.left.IntSlice[0]) / 2
}

type hp struct{ sort.IntSlice }

func (h *hp) Push(v any) {
 h.IntSlice = append(h.IntSlice, v.(int))
}

func (h *hp) Pop() any {
 a := h.IntSlice
 v := a[len(a)-1]
 h.IntSlice = a[:len(a)-1]
 return v
}

好啦，这次的笔记就先到这。中位数看着像要排序，其实只要把两边的边界守好，答案就一直稳稳待在堆顶啦。