使用堆和优先队列求解数据流中的中位数

剑指 Offer 41. 数据流中的中位数

题目

• 从一个数据流中求解当前已读数据的中位数。
• 如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。

题解

• 维护两个堆，大根堆LargeRootHeap中保存所有比当前中位数小的数，小根堆SmallRootHeap中保存所有比当前中位数大的数。所以，这个两堆的堆顶元素一定是距离中位数最近的两个数。
  
  • 根据上图可以看到，只要满足堆的定义即可，不需要满足二叉搜索树。之前一直纠结是不是要排序，但是对于大根堆，只要维护最大的数在堆顶；对于小根堆，维护最小的数在堆顶即可。
• 对堆进行增加元素和删除元素时，要基于队列操作，即只能从队尾插入元素，从队头删除元素，同时要保持堆性质不变。
  
  
• 当LargeRootHeap和SmallRootHeap长度一样时，将元素插入SmallRootHeap；否则，插入LargeRootHeap。
• 插入规则：
  • 为了满足SmallRootHeap中的元素都比中位数大，读取的当前元素入队LargeRootHeap，将LargeRootHeap的堆顶元素(即LargeRootHeap中最大值)入队SmallRootHeap。
  • LargeRootHeap同理。

• 根据以上策略可以维持：$0<=SmallRootHeap.size-LargeRootHeap.size<=1$，对于奇数个元素，中位数是：$SmallRootHeap.head$；偶数个元素，中位数是：$(SmallRootHeap.head+LargeRootHeap.head)/2$
  

JS代码实现

实现堆队列

• 由于JS内部没有封装相关的数据结构，这块得自己写。

// 基类，实现二叉树的相关方法
class BasicBinaryTree {
  constructor (list = []) {
    this.list = list
  }
  size () {
    return this.list.length
  }
  parent (i) {
    let index = -1
    if (i & 1) {
      index = (i - 1) / 2
    } else {
      index = (i - 2) / 2
    }
    return index >= 0 ? index : -1 
  }
  leftChild (i) {
    return 2 * i + 1
  }
  rightChild (i) {
    return 2 * i + 2
  }
  isEmpty () {
    return this.size() === 0
  }
  head () {
    return this.isEmpty() ? null : this.list[0]
  }
  isGreater (i, j) {
    if (this.get(i) > this.get(j)) {
      return true
    }
    return false
  }
  swap(i, j) {
    const v = this.get(i)
    this.set(i, this.get(j))
    this.set(j, v)
  }
  set (i, e) {
      this.list[i] = e
  }
  get (i) {
    if (i < 0 || i >= this.size()) {
      return null
    }
    return this.list[i]
  }
}

// 大根堆
class LargeRootHeap extends BasicBinaryTree {
  // 递归
  heapifyDownRecursion (i, size = this.size) {
    const rightChild = this.rightChild(i)
    const leftChild = this.leftChild(i)
    let swapIndex = this.isGreater(rightChild, leftChild) ? rightChild : leftChild
    if (swapIndex < size && this.isGreater(swapIndex, i)) {
      this.swap(swapIndex, i)
      this.heapifyDownRecursion(swapIndex, size)
    }
  }

   // 非递归
   heapifyDown (i, size = this.size()) {
    while (i < size) {
      const rightChild = this.rightChild(i)
      const leftChild = this.leftChild(i)
      let swapIndex = -1
      if (rightChild < size && leftChild < size) {
          swapIndex = this.isGreater(rightChild, leftChild) ? rightChild : leftChild
      } else if (rightChild >= size) {
          swapIndex = leftChild
      } else {
          swapIndex = rightChild
      }
      if (swapIndex >= size || !this.isGreater(swapIndex, i)) {
        return
      }
      this.swap(swapIndex, i)
      i = swapIndex
    }
  }
  
  heapifyUp (i) {
    const parent = this.parent(i)
    if (i >= 0 && parent >=0 && this.isGreater(i, parent)){
      this.swap(i, parent)
      this.heapifyUp(parent)
    }
  }
}

// 大根堆队列
class LargeRootQueue extends LargeRootHeap {
  enQueue (e) {
    this.list.push(e)
    this.heapifyUp(this.size() - 1)
  }
  deQueue () {
    const head = this.head()
    if (head !== null) {
      const last = this.list.pop()
      if (head !== last) {
        this.list[0] = last
        this.heapifyDown(0, this.size())
      }
    }
    return head
  }
}

// 小根堆队列
class SmallRootQueue extends LargeRootQueue {
  // 小根堆只是每次比较并操作最小的数，只修改isGreater中的判断逻辑，其余直接复用即可
  isGreater (a, b) {
    if (this.get(a) < this.get(b)) {
      return true
    }
    return false
  }
}

从数据流中获取中位数

/**
 * initialize your data structure here.
 */
var MedianFinder = function() {
  this.largeRootQueue = new LargeRootQueue()  // 大根堆，保存比中位数小的元素
  this.smallRootQueue = new SmallRootQueue() // 小根堆，保存比中位数大的元素
};

/** 
 * @param {number} num
 * @return {void}
 * 说明：
 * - 如果largeSize === smallSize，将元素插入largeHeap中，largeHeap堆顶出队，插入smallHeap中
 * - 否则，将元素插入smallHeap中，smallHeap堆顶出队，插入largeHeap中
 */
MedianFinder.prototype.addNum = function(num) {
  const largeRootQueueSize = this.largeRootQueue.size()
  const smallRootQueueSize = this.smallRootQueue.size()
  if (largeRootQueueSize === smallRootQueueSize) { // 向small中添加元素
    this.largeRootQueue.enQueue(num)
    this.smallRootQueue.enQueue(this.largeRootQueue.deQueue())
  } else { // 向large中添加元素
    this.smallRootQueue.enQueue(num)
    this.largeRootQueue.enQueue(this.smallRootQueue.deQueue())
  }
};

/**
 * @return {number}
 * 说明：
 * - 当已读奇数个元素，即largeSize === smallSize，中位数 = smallHeap堆顶元素
 * - 否则，中位数 = (smallHeap堆顶元素 + largeHeap堆顶元素) / 2
 */
MedianFinder.prototype.findMedian = function() {
  if (this.largeRootQueue.size() === this.smallRootQueue.size()) {
    return (this.largeRootQueue.head() + this.smallRootQueue.head()) / 2
  }
  return this.smallRootQueue.head()
};