← 中文资源

简体中文版本

双堆中位数追踪

用大顶堆保存较小一半、小顶堆保存较大一半,在对数时间内维护数据流的中位数。

阅读英文原文

功能说明

数据流中位数无法通过排序整批数据得到,因为新元素会持续到达。双堆方案把数据分成两半:左侧大顶堆保存较小的一半,堆顶是这一半的最大值;右侧小顶堆保存较大的一半,堆顶是这一半的最小值。两侧长度差不超过一,中位数只由两个堆顶决定。

每次插入只需要对数时间,查询中位数是常数时间。示例依次插入七、三、九、一、五,并打印每一步的中位数和两堆大小,可以看到奇数与偶数长度的切换。

关键代码

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())
	}
}