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

package main

import (
	"math"
)

// 大根堆小根堆
// 最终取中间的两个元素作为中位数或者取中间一个
// 左边取最大的，因此使用排序后是递增序列的大根堆
// 右边取最小的，因此使用排序后是递减序列的小根堆
// 需要保证所有右侧元素都比左侧堆元素大
// 为了保证两侧元素均匀分布，我们可以设置，如果当前堆总长度为奇数时插入左侧大根堆，否则插入右侧
// 如果插入某个元素时，它原本应该分配到右边，但是它的值却没有左侧大根堆堆顶元素大怎么办？
// 先将左侧堆顶元素移除压入右侧小根堆，然后将新元素压入大根堆
// 同理，将右侧小根堆堆顶元素移除，插入到左侧，然后将新元素插入到右侧
// 首先，我们需要先实现了大根堆和小根堆（有序队列形式的（支持 Push 和 Pop 操作））
type smallHeap struct {
	Size int
	Array []int
}

func (sh *smallHeap) Push(x int) {
	sh.Array = append(sh.Array, x)
	// 插入索引
	index := sh.Size
	sh.Size++

	// 向上调整
	for index >= 0 {
		root := (index-1)/2

		if sh.Array[root] > sh.Array[index] {
			sh.Array[root], sh.Array[index] = sh.Array[index], sh.Array[root]
			index = root
		} else {
			break
		}
	}
}

func (sh *smallHeap) Pop() int {
	if sh.Size == 0 {
		return math.MinInt
	}
	// 弹出堆顶元素
	x := sh.Array[0]
	// 首尾交换
	sh.Array[0], sh.Array[sh.Size-1] = sh.Array[sh.Size-1], sh.Array[0]
	// 弹出
	sh.Array = sh.Array[:sh.Size-1]
	sh.Size--
	// 向下调整
	root := 0

	for root < sh.Size {
		left, right := 2*root+1, 2*root+2
		if left >= sh.Size {
			break
		}

		if right < sh.Size && sh.Array[right] < sh.Array[left] {
			left = right
		}

		if sh.Array[left] > sh.Array[root] {
			break
		}

		sh.Array[root], sh.Array[left] = sh.Array[left], sh.Array[root]
		root = left
	}

	return x
}


func (sh *smallHeap) Top() int {
	if sh.Size == 0 {
		return math.MinInt
	}

	return sh.Array[0]
}


type bigHeap struct {
	Size int
	Array []int
}

func (bh *bigHeap) Push(x int) {
	bh.Array = append(bh.Array, x)
	// 插入索引
	index := bh.Size
	bh.Size++

	// 向上调整
	for index >= 0 {
		root := (index-1)/2

		if bh.Array[root] < bh.Array[index] {
			bh.Array[root], bh.Array[index] = bh.Array[index], bh.Array[root]
			index = root
		} else {
			break
		}
	}
}

func (bh *bigHeap) Pop() int {
	if bh.Size == 0 {
		return math.MinInt
	}
	// 弹出堆顶元素
	x := bh.Array[0]
	// 首尾交换
	bh.Array[0], bh.Array[bh.Size-1] = bh.Array[bh.Size-1], bh.Array[0]
	// 弹出
	bh.Array = bh.Array[:bh.Size-1]
	bh.Size--
	// 向下调整
	root := 0

	for root < bh.Size {
		left, right := 2*root+1, 2*root+2
		if left >= bh.Size {
			break
		}

		if right < bh.Size && bh.Array[right] > bh.Array[left] {
			left = right
		}

		if bh.Array[left] < bh.Array[root] {
			break
		}

		bh.Array[root], bh.Array[left] = bh.Array[left], bh.Array[root]
		root = left
	}
	
	return x
}

func (bh *bigHeap) Top() int {
	if bh.Size == 0 {
		return math.MinInt
	}

	return bh.Array[0]
}

var small smallHeap
var big bigHeap

// 左大右小
func Insert(num int){
	size := small.Size + big.Size
	if size % 2 == 0 {
		if small.Size > 0 && num > small.Top() {
			big.Push(small.Pop())
			small.Push(num)
		} else {
			big.Push(num)
		}
	} else {
		if big.Size > 0 && num < big.Top() {
			small.Push(big.Pop())
			big.Push(num)
		} else {
			small.Push(num)
		}
	}
}

func GetMedian() float64{
	size := small.Size + big.Size
	if size % 2 == 1 {
		return float64(big.Top())
	} else {
		return float64(small.Top()+big.Top())/2
	}
}