堆-数据流的中位数

72 阅读1分钟

295. 数据流的中位数

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

例如,

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

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

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

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

示例:

addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3) 
findMedian() -> 2

进阶:

  1. 如果数据流中所有整数都在 0 到 100 范围内,你将如何优化你的算法?
  2. 如果数据流中 99% 的整数都在 0 到 100 范围内,你将如何优化你的算法?
class MedianFinder {
  constructor() {
    this.size = 0;
    this.leftHeap = new Heap((a, b) => a > b); // 左边大根堆
    this.rightHeap = new Heap((a, b) => a < b); // 右边小根堆
  }

  addNum(num) {
    const { leftHeap, rightHeap } = this;
    if (this.size % 2 === 0) {
      if (!rightHeap.getTop() || num <= rightHeap.getTop()) {
        leftHeap.push(num);
      } else {
        const rightTop = rightHeap.pop();
        rightHeap.push(num);
        leftHeap.push(rightTop);
      }
    } else {
      if (num >= leftHeap.getTop()) {
        rightHeap.push(num);
      } else {
        const leftTop = leftHeap.pop();
        leftHeap.push(num);
        rightHeap.push(leftTop);
      }
    }

    this.size++;
  }

  findMedian() {
    if (this.size % 2 === 0) {
      return (this.leftHeap.getTop() + this.rightHeap.getTop()) / 2;
    } else {
      return this.leftHeap.getTop();
    }
  }
}

class Heap {
  constructor(compare) {
    this.nodeList = [];
    this.compare = typeof compare === 'function' ? compare : this.defaultCompare;
  }

  defaultCompare(a, b) {
    return a > b;
  }

  getSize() {
    return this.nodeList.length;
  }

  getTop() {
    if (this.nodeList.length === 0) {
      return null;
    }

    return this.nodeList[0];
  }

  push(value) {
    this.nodeList.push(value);
    this.up(this.nodeList.length - 1);
  }

  up(index) {
    const { compare, parent, nodeList } = this;
    let curIndex = index;
    let parentIndex = parent(curIndex);
    while (curIndex > 0 && compare(nodeList[curIndex], nodeList[parentIndex])) {
      const temp = nodeList[curIndex];
      nodeList[curIndex] = nodeList[parentIndex];
      nodeList[parentIndex] = temp;
      curIndex = parentIndex;
      parentIndex = parent(curIndex);
    }
  }

  parent(index) {
    return index % 2 === 0 ? index / 2 - 1 : (index - 1) / 2;
  }

  pop() {
    if (this.nodeList.length === 0) {
      return null;
    }

    const top = this.nodeList[0];
    this.nodeList[0] = this.nodeList[this.nodeList.length - 1];
    this.nodeList.pop();
    this.down(0);
    return top;
  }
  
  down(index) {
    const { compare, left, right, nodeList } = this;
    let curIndex = index;
    while (left(curIndex) < nodeList.length) {
      let target = left(curIndex);
      if (right(curIndex) < nodeList.length && compare(nodeList[right(curIndex)], nodeList[target])) {
        target = right(curIndex);
      }

      if (compare(nodeList[curIndex], nodeList[target])) {
        break;
      }

      const temp = nodeList[curIndex];
      nodeList[curIndex] = nodeList[target];
      nodeList[target] = temp;
      curIndex = target;
    }
  }

  left(index) {
    return index * 2 + 1;
  }

  right(index) {
    return index * 2 + 2;
  }

}