堆的使用：如何在大数据流中快速找到中位数

这个问题从一个leetcode算法题说起(295.数据流中位数)。

求一个无序数组的中位数，我们首先想到的是进行排序，然后找到中间值（如果数组长度是奇数，那么正好是中间那个值；如果数组长度是偶数，那么是中间两个数的平均值）。但是我们只是找到中间值而已，不需要一个完全排序的数组，只要保证数组的中间值前面的部分的值都是小于中间值，而数组的中间值后面的部分的值大于中间值，就可以正确找到中间值了。

正如标题所写的，使用堆的数据结构就能以O(logn)时间复杂度来解决这个问题。

堆(Heap)数据结构介绍

堆结构是一个完全二叉树，堆结构分为小顶堆和大顶堆。以小顶堆为例，小顶堆中，二叉树的根节点的值，小于左右两个节点的值。

由于堆是一个完全二叉树，我们可以通过一个数组来存储数据。

我们将数组的第一个小标冗余出来，方便计算。从小标为1开始存储。root_index代表根节点下标，left_index代表左节点下标，right_index代表右节点下标。计算公式为：

left_index = root_index * 2

right_index = root_index * 2 + 1

解题思路

1. 创建一个大顶堆，存放小值集合；创建一个小顶堆，存放大值集合
2. 保持大小两个顶堆平衡，小顶堆的数量比大顶堆数量不超过1
3. 最后，如果大小顶堆数量一样，就取大小顶堆顶值的平均值，如果不一样，就取小顶堆顶值

手撕一个堆操作

type MedianHeap struct {
    List []int  //数据存储
    Sort int8 //1为大顶堆，-1为小顶堆
    Size int  //存储量
}

// 创建堆
func NewMedianHeap(sort int8) *MedianHeap {
    return &MedianHeap{
       List: []int{0}, //为方便计算index，第一个值赘余
       Sort: sort,
    }
}

// 获取堆的顶值
func (h *MedianHeap) GetTop() int {
    if h.Size > 0 {
       return h.List[1]
    }

    if h.Sort == 1 {
       return -1 << 31
    } else {
       return 1 << 31
    }
}

// 替换顶值
func (h *MedianHeap) ReplaceTop(val int) (top int) {
    if h.Size > 0 {
       top = h.List[1]
    }
    h.List[1] = val
    h.down()
    return
}

// 向尾部追加值
func (h *MedianHeap) ToTail(val int) {
    h.List = append(h.List, val)
    h.Size++
    h.up()
    return
}

//向下交换
func (h *MedianHeap) down() {
    var (
       parent      = 1
       left, right int
    )
    left = parent * 2
    right = parent*2 + 1

    for right <= h.Size {
       var compare int
       if h.Sort == 1 {
          if h.List[left] > h.List[right] {
             compare = left
          } else {
             compare = right
          }

          if h.List[compare] > h.List[parent] {
             h.swap(compare, parent)
             parent = compare
          } else {
             break
          }
       }

       if h.Sort == -1 {
          if h.List[left] < h.List[right] {
             compare = left
          } else {
             compare = right
          }

          if h.List[compare] < h.List[parent] {
             h.swap(compare, parent)
             parent = compare
          } else {
             break
          }
       }

       left = parent * 2
       right = parent*2 + 1
    }

    if left == h.Size {
       if h.Sort == 1 && h.List[left] > h.List[parent] {
          h.swap(left, parent)
       }

       if h.Sort == -1 && h.List[left] < h.List[parent] {
          h.swap(left, parent)
       }
    }
}

//向上交换
func (h *MedianHeap) up() {
    compare := h.Size
    parent := compare / 2
    for parent > 0 {
       if h.Sort == 1 {
          if h.List[compare] > h.List[parent] {
             h.swap(compare, parent)
             compare = parent
          } else {
             break
          }
       }

       if h.Sort == -1 {
          if h.List[compare] < h.List[parent] {
             h.swap(compare, parent)
             compare = parent
          } else {
             break
          }
       }

       parent = compare / 2
    }
}

func (h *MedianHeap) swap(i, j int) {
    h.List[i], h.List[j] = h.List[j], h.List[i]
}

求数据流的中位数

type MedianFinder struct {
    bigHeap   *MedianHeap //大顶堆
    smallHeap *MedianHeap //小顶堆
}

func Constructor() MedianFinder {
    return MedianFinder{
       bigHeap:   NewMedianHeap(1),
       smallHeap: NewMedianHeap(-1),
    }
}

func (this *MedianFinder) AddNum(num int) {
    if this.bigHeap.Size == this.smallHeap.Size {
       if num >= this.bigHeap.GetTop() {
          this.smallHeap.ToTail(num)
       } else {
          top := this.bigHeap.ReplaceTop(num)
          this.smallHeap.ToTail(top)
       }
    } else {
       if num <= this.smallHeap.GetTop() {
          this.bigHeap.ToTail(num)
       } else {
          top := this.smallHeap.ReplaceTop(num)
          this.bigHeap.ToTail(top)
       }
    }
}

func (this *MedianFinder) FindMedian() float64 {
    if this.bigHeap.Size == this.smallHeap.Size {
       return float64(this.bigHeap.GetTop()+this.smallHeap.GetTop()) / 2
    } else {
       return float64(this.smallHeap.GetTop())
    }
}