功能说明
数据流中位数无法通过排序整批数据得到,因为新元素会持续到达。双堆方案把数据分成两半:左侧大顶堆保存较小的一半,堆顶是这一半的最大值;右侧小顶堆保存较大的一半,堆顶是这一半的最小值。两侧长度差不超过一,中位数只由两个堆顶决定。
每次插入只需要对数时间,查询中位数是常数时间。示例依次插入七、三、九、一、五,并打印每一步的中位数和两堆大小,可以看到奇数与偶数长度的切换。
关键代码
maxHeap 的 Less 使用大于号,使 container/heap 按大顶堆工作;minHeap 使用小于号。Push 和 Pop 必须定义在指针接收者上,类型断言把 any 转回 int。Add 根据左堆堆顶决定插入哪一侧,再调用 rebalance:左堆比右堆多两个时把一个元素搬到右侧,右堆比左堆多时搬回左侧。Median 在两侧相等时取两个堆顶的平均值,否则取左堆堆顶。
使用注意
堆中只保存整数,若需要浮点数应把类型改为 float64 并调整断言。堆不支持删除任意元素,滑动窗口中位数需要额外维护延迟删除表。并发写入要加锁。整数相加可能溢出,样本范围大时先转 float64 再求和。两个堆都为空时返回零,调用方应区分零与无数据,或增加一个计数方法。
验证方式
插入七后中位数为七;插入三后为五;插入九后为七;插入一后为五;插入五后为五。每一步左堆大小或等于右堆,或比右堆多一,输出中的 lower 与 upper 计数应满足该不变式。可以打乱插入顺序,最终中位数仍应为五。
代码示例
package main
import (
"container/heap"
"fmt"
)
type maxHeap []int
func (h maxHeap) Len() int { return len(h) }
func (h maxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h maxHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *maxHeap) Push(x any) { *h = append(*h, x.(int)) }
func (h *maxHeap) Pop() any {
old := *h
n := len(old)
value := old[n-1]
*h = old[:n-1]
return value
}
type minHeap []int
func (h minHeap) Len() int { return len(h) }
func (h minHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h minHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *minHeap) Push(x any) { *h = append(*h, x.(int)) }
func (h *minHeap) Pop() any {
old := *h
n := len(old)
value := old[n-1]
*h = old[:n-1]
return value
}
type MedianTracker struct {
lower maxHeap
upper minHeap
}
func (t *MedianTracker) Add(value int) {
if t.lower.Len() == 0 || value <= t.lower[0] {
heap.Push(&t.lower, value)
} else {
heap.Push(&t.upper, value)
}
t.rebalance()
}
func (t *MedianTracker) rebalance() {
for t.lower.Len() > t.upper.Len()+1 {
heap.Push(&t.upper, heap.Pop(&t.lower))
}
for t.upper.Len() > t.lower.Len() {
heap.Push(&t.lower, heap.Pop(&t.upper))
}
}
func (t *MedianTracker) Median() float64 {
if t.lower.Len() == 0 {
return 0
}
if t.lower.Len() == t.upper.Len() {
return float64(t.lower[0]+t.upper[0]) / 2
}
return float64(t.lower[0])
}
func main() {
tracker := &MedianTracker{}
for _, value := range []int{7, 3, 9, 1, 5} {
tracker.Add(value)
fmt.Printf("after %d: median=%.1f lower=%d upper=%d\n", value, tracker.Median(), tracker.lower.Len(), tracker.upper.Len())
}
}