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

package main

// 手写实现小顶堆
type smallTopHeap struct {
	Size  int
	Array []int
}

// 手写实现大顶堆
type bigTopHeap struct {
	Size  int
	Array []int
}

// 小顶堆的两个核心方法之 Push
func (h *smallTopHeap) Push(x int) {
	h.Array = append(h.Array, x)
  	// 往堆中添加一个新的元素，对应下标刚好为 h.Size，需要判断这个元素所在分支是否满足小顶堆性质
  	// 如果不满足，与父节点交换，一直向上调整
	smallup(h.Array, h.Size)
  	// 小顶堆中元素数量+1
	h.Size++
}
// 小顶堆插入元素时的向上调整
func smallup(arr []int, index int) {
  	// index 是孩子节点的下标
	for index > 0 {
	  	// 比如孩子节点下标为 3，它的父节点下标为 1 即（index-1）/ 2
		root := (index - 1) / 2
	  	// 如果根节点的值大于孩子节点的值，那么为了满足小顶堆性质，需要交换
		if arr[root] > arr[index] {
			arr[index], arr[root] = arr[root], arr[index]
			index = root // 更新孩子的最新下标，继续往上调整
		} else { // 如果符合堆性质，直接退出
			break
		}
	}
}

// 小顶堆的两个核心方法之 Pop
func (h *smallTopHeap) Pop() int {
	// 记录弹出的最小元素值
	x := h.Array[0]
	// 交换首尾元素，最小元素值交换到末尾
	h.Array[0], h.Array[h.Size-1] = h.Array[h.Size-1], h.Array[0]
  	// 将最小元素值从数组中移除（弹出）
	h.Array = h.Array[:h.Size-1]
  	// 更新堆的长度
	h.Size--
  	// 我们将堆末尾元素交换到根节点，需要从根节点向下调整使其重复符合小顶堆性质
  	// 其中 0 是代表我们想要调整元素的下标，h.Size 代表整个堆的长度
	smalldown(h.Array, 0, h.Size)
	return x
}

// 弹出元素时，需要对交换到根节点的元素进行向下调整
func smalldown(arr []int, index, end int) {
  	// end 是我们堆长度，下标不能超过 end-1 
	for index < end {
	  	// 两个孩子下标
		left, right := 2*index+1, 2*index+2
	  	// 如果左孩子的下标都超出范围，说明已经遍历到叶子节点了，可以退出
		if left >= end {
			break
		}
		// 如果存在右孩子，并且右孩子值比左孩子值更小，那么我们把右孩子的下标赋给左孩子的下标（统一使用 left 与父节点值比较）
		if right < end && arr[right] < arr[left] {
			left = right
		}
	  	// 如果父节点的值比两个孩子中最小的值还小（符合小顶堆性质），那么直接退出
		if arr[index] < arr[left] {
			break
		}
		// 将较小的元素值交换到父节点位置
		arr[index], arr[left] = arr[left], arr[index]
	  	// 更新需要调整的下标，继续往下调整
		index = left
	}
}

// 获取堆顶的值（小根堆堆顶是数组中的最小值，它位于数组的第一个位置）
// （大根堆堆顶是数组中的最大值，它位于大根堆数组的第一个位置）
// 经过堆排序后，小根堆最后得到一个递增的序列；大根堆得到一个递减的序列；
func (h smallTopHeap) Top() int {
  	// 如果堆为空，返回一个 32 位的最小负数进行标识
	if h.Size == 0 {
		return 1 << 32 - 1
	}
	// 返回堆顶元素
	return h.Array[0]
}

// 大顶堆的实现（和小顶堆实现思路类似）
func (h *bigTopHeap) Push(x int) {
	h.Array = append(h.Array, x)
	h.Size++
	bigup(h.Array, h.Size-1)
}

func bigup(arr []int, index int) {
	for index > 0 {
		root := (index - 1) / 2
		if arr[root] < arr[index] {
			arr[index], arr[root] = arr[root], arr[index]
			index = root
		} else {
			break
		}
	}
}

func (h *bigTopHeap) Pop() int {
	x := h.Array[0]

	h.Array[0], h.Array[h.Size-1] = h.Array[h.Size-1], h.Array[0]
	h.Array = h.Array[:h.Size-1]
	h.Size--
	bigdown(h.Array, 0, h.Size)
	return x
}

func bigdown(arr []int, index, end int) {
	for index < end {
		left, right := 2*index+1, 2*index+2
		if left >= end {
			break
		}

		if right < end && arr[right] > arr[left] {
			left = right
		}
		if arr[index] > arr[left] {
			break
		}

		arr[index], arr[left] = arr[left], arr[index]
		index = left
	}
}

func (h bigTopHeap) Top() int {
	if h.Size == 0 {
		return 1 << 32 - 1
	}

	return h.Array[0]
}

// 这一题出的感觉有点怪，我们需要申请全局变量
var min = &smallTopHeap{}
var max = &bigTopHeap{}

// 这个函数就是从输入流中读取一个新的值插入到某个堆中
// 我们要求的是中位值，不妨使用两个堆，一个小顶堆（递减序列），一个大顶堆（递增序列）
// 中位值左边元素均小于它，右边元素均大于它
// 当总长度为偶数时，存在两个中位值分割数组；当总长度为奇数时，只存在一个中位值；
// 我们的实现的思路就是：使用两个堆来插入数组的元素，保证左边堆的元素均小于右边堆的元素，而且最好实现 O(1) 时间内取得最终的中位值。这就需要确保中位值要么处于大顶堆的堆顶，要么处于小顶堆的堆顶。所以我们使用大顶堆存放左侧小的那部分元素，这样就保证堆顶元素是左边一半的最大值；然后我们再使用小顶堆存放右侧大的那部分元素，这样我们又可以保证堆顶元素是右侧那部分元素的最小值；
// 最终计算结果时，如果两部分总长度为偶数，分别取两个堆顶元素进行相加取平均值就是中位值，而如果两部分总长度为奇数，我们取长度较长的那个堆的堆顶元素就是中位值。

// 为了均匀分配到两个堆中，规定：当当前两个堆的总长度为偶数时，优先插入小根堆，否则优先插入大根堆
// 左边（大根堆，递增）右边（小顶堆，递减）
// 特殊情况处理：如果当前两个堆的总长度为偶数，优先插入小根堆，我们之前的设计是：小根堆存放右边比较大的那部分元素，假设此时，待插入的值比左边大顶堆的堆顶元素还小，该怎么处理呢？如果将这个元素插入小根堆就不符合右侧全部元素大于左侧了，处理方法：先将左边大顶堆的堆顶元素弹出压入到小顶堆中，然后将新的较小元素再压入小根堆中。
// 同理，当两个堆的总长度为奇数，优先插入大根堆，但是当前插入元素比小根堆的堆顶元素还要大，那么我们为了保证左边小于右边的性质，先将小顶堆堆顶元素弹出压入大顶堆中，然后将这个比较大的元素再压入小顶堆即可。
// 非特殊情况，按照优先级进行处理
func Insert(num int) {
	size := min.Size + max.Size
	if size%2 == 0 { 
		if max.Size > 0 && num < max.Top() { // 只有另一个堆中有元素时才进行特殊情况判断
			min.Push(max.Pop())
			max.Push(num)
		} else {
			min.Push(num)
		}
	} else {
		if min.Size > 0 && num > min.Top() {
			max.Push(min.Pop())
			min.Push(num)
		} else {
			max.Push(num)
		}
	}
}

// 两个堆最初状态是总长度为偶数，因此优先插入的是小顶堆（右侧），所以当总长度为奇数时，取的是小顶堆的堆顶元素作为中位值
func GetMedian() float64 {
	if (min.Size+max.Size)%2 == 1 {
		return float64(min.Top())
	} else {
		return float64(min.Top()+max.Top()) / 2
	}
}