[路飞] 295.数据流的中位数

题目描述

中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。

例如，

[2,3,4] 的中位数是 3

[2,3] 的中位数是 (2 + 3) / 2 = 2.5

设计一个支持以下两种操作的数据结构：

void addNum(int num) - 从数据流中添加一个整数到数据结构中。 double findMedian() - 返回目前所有元素的中位数。

解题思路

通过维护对顶堆，保证前面大顶堆的个数不超过小顶堆的元素个数 + 1.

最终返回两个堆顶元素的平均数或者大顶堆的堆顶元素

元素放入过程

如果大顶堆是空，直接往里放，或者元素小于等于大顶堆堆顶元素，就直接往里放；其余情况都往右边小顶堆放。

如果小顶堆元素个数大于大顶堆，从小顶堆中 extract 一个元素到大顶堆

如果大顶堆元素个数 === 小顶堆 + 2，从大顶堆中 extract 一个元素刀小顶堆

找中位数

根据大顶堆元素个数是否和小顶堆相等，返回两个堆顶元素相加除 2，或返回左边大顶堆堆顶元素

代码

var MedianFinder = function () {
  this.l = new Heap((a, b) => a - b > 0)
  this.r = new Heap((a, b) => a - b < 0)
}

/**
 * @param {number} num
 * @return {void}
 */
MedianFinder.prototype.addNum = function (num) {
  if (!this.l.size() || this.l.top() <= num) {
    this.l.insert(num)
  } else {
    this.r.insert(num)
  }

  if (this.r.size() > this.l.size()) {
    this.l.insert(this.r.extract())
  }

  if (this.l.size() === this.r.size() + 2) {
    this.r.insert(this.l.extract())
  }
}

/**
 * @return {number}
 */
MedianFinder.prototype.findMedian = function () {
  if (this.l.size() === this.r.size()) {
    return (this.l.top() + this.r.top()) / 2
  }

  return this.l.top()
}

class Heap {
  constructor(compareFn) {
    this.heap = []
    this.compareFn = compareFn
  }

  getLeftIndex(index) {
    return index * 2 + 1
  }

  getRightIndex(index) {
    return index * 2 + 2
  }

  getParentIndex(index) {
    return Math.floor((index - 1) / 2)
  }

  size() {
    return this.heap.length
  }

  isEmpty() {
    return this.size() === 0
  }

  swap(parent, index) {
    const arr = this.heap
    ;[arr[parent], arr[index]] = [arr[index], arr[parent]]
  }

  insert(value) {
    const index = this.heap.length
    this.heap.push(value)

    this.siftUp(index)
  }

  siftUp(index) {
    let parent = this.getParentIndex(index)

    while (index > 0 && this.compareFn(this.heap[parent], this.heap[index])) {
      this.swap(parent, index)
      index = parent
      parent = this.getParentIndex(index)
    }
  }

  extract() {
    if (this.isEmpty()) return
    if (this.size() === 1) return this.heap.pop()

    const removedItem = this.heap[0]
    this.heap[0] = this.heap.pop()
    this.siftDown(0)

    return removedItem
  }

  siftDown(index) {
    let element = index
    const left = this.getLeftIndex(index)
    const right = this.getRightIndex(index)

    if (
      index < this.size() &&
      this.compareFn(this.heap[element], this.heap[left])
    ) {
      element = left
    }

    if (
      index < this.size() &&
      this.compareFn(this.heap[element], this.heap[right])
    ) {
      element = right
    }

    if (element !== index) {
      this.swap(element, index)
      this.siftDown(element)
    }
  }

  top() {
    return this.heap[0]
  }
}