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

package main

import (
	"container/heap"
)

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

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

type IntHeap []int
type InvertIntHeap []int

// 使用全局变量来维护两个堆
var minHeap *IntHeap
var maxHeap *InvertIntHeap

func init() {
	// 在init函数中初始化堆
	minH := &IntHeap{}
	maxH := &InvertIntHeap{}
	heap.Init(minH)
	heap.Init(maxH)
	minHeap = minH
	maxHeap = maxH
}

func (h InvertIntHeap) Len() int           { return len(h) }
func (h InvertIntHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h InvertIntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *InvertIntHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

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

func Insert(num int) {
	if maxHeap.Len() == 0 || num <= (*maxHeap)[0] {
		heap.Push(maxHeap, num)
	} else {
		heap.Push(minHeap, num)
	}

	// Adjust the heaps to make sure their size difference is no more than 1
	if maxHeap.Len() > minHeap.Len()+1 {
		heap.Push(minHeap, heap.Pop(maxHeap))
	} else if minHeap.Len() > maxHeap.Len()+1 {
		heap.Push(maxHeap, heap.Pop(minHeap))
	}
}

func GetMedian() float64 {
	if maxHeap.Len() > minHeap.Len() {
		return float64((*maxHeap)[0])
	} else if minHeap.Len() > maxHeap.Len() {
		return float64((*minHeap)[0])
	}
	return float64((*maxHeap)[0]+(*minHeap)[0]) / 2.0
}