题解 | #数据流中的中位数#

package main

import "container/heap"

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 interface{}) {
	// Push and Pop use pointer receivers because they modify the slice's length,
	// not just its contents.
	*h = append(*h, x.(int))
}

func (h *MaxHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

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 interface{}) {
	// Push and Pop use pointer receivers because they modify the slice's length,
	// not just its contents.
	*h = append(*h, x.(int))
}

func (h *MinHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// 中位数的特征，它是数组中间个数字或者两个数字的均值，它是数组较小的一半元素中最大的一个，同时也是数组较大的一半元素中最小的一个
// 那我们只要每次维护最小的一半元素和最大的一半元素，并能快速得到它们的最大值和最小值
var (
	maxNumbers MinHeap
	minNumbers MaxHeap
)

func Insert(num int) {
	// 小顶堆，用于存储较大的值，其中顶部最小
	heap.Init(&maxNumbers)
	// 大顶堆，用于存储较小的值，其中顶部最大
	heap.Init(&minNumbers)
	// 每次输入的数据流先进入大顶堆排序，然后将大顶堆的最大值弹入小顶堆中，完成整个的排序
	heap.Push(&minNumbers, num)
	heap.Push(&maxNumbers, heap.Pop(&minNumbers))
	// 因为大顶堆的数据不可能会比小顶堆少一个
	// 因此需要再比较二者的长度，若是大顶堆长度小于小顶堆，需要从小顶堆中弹出最小值到大顶堆中进行平衡
	if len(minNumbers) < len(maxNumbers) {
		heap.Push(&minNumbers, heap.Pop(&maxNumbers))
	}
}

func GetMedian() float64 {
	// 中位数只会在两个堆的堆顶出现
	// 约定奇数个元素时取大顶堆的顶部值，偶数个元素时取两堆顶的平均值
	if len(minNumbers) > len(maxNumbers) {
		return float64(minNumbers[0])
	}

	return (float64(minNumbers[0]) + float64(maxNumbers[0])) / 2
}