Hard | LeetCode 295 | 剑指 Offer 41. 数据流中的中位数 | 优先队列(堆排序)

Hard | LeetCode 295 | 剑指 Offer 41. 数据流中的中位数 | 优先队列(堆排序)

如何得到一个数据流中的中位数?如果从数据流中读出奇数个数值,那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值,那么中位数就是所有数值排序之后中间两个数的平均值。

例如,

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

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

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

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

示例 1:

输入:
["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]
[[],[1],[2],[],[3],[]]
输出:[null,null,null,1.50000,null,2.00000]

示例 2:

输入:
["MedianFinder","addNum","findMedian","addNum","findMedian"]
[[],[2],[],[3],[]]
输出:[null,null,2.00000,null,2.50000]

限制:

  • 最多会对 addNum、findMedian 进行 50000 次调用。

方法一 : 优先队列

又是一道非常巧妙的Hard题。它的思路是把数据元素平均分成两部分存储, 左半部分使用大顶堆存储, 右半部分使用小顶堆存储, 并且大顶堆的所有元素, 小于小顶堆的所有元素, 这样整个的中位数就可以通过两个堆的堆顶元素获得了。

那么问题来了, 在添加元素的过程, 如何保证两个堆的元素的数量相等呢?

一个很朴素的想法: 添加元素时, 轮流向大顶堆和小顶堆添加元素, 这样两个堆的数量相差值不会超过1。

那么问题又来了, 假如我现在要添加一个元素到大顶推, 如果要添加到大顶堆的元素比小顶堆的一些元素大怎么办, 我直接插进去不就不能保证大顶堆元素小于小顶堆了?

这个解决办法非常巧妙: 首先把当前的元素先插到小顶堆, 然后把小顶堆的堆顶元素插入到大顶堆。就完美解决上述问题了。因为(1) 加入当前元素比某些小顶堆元素大, 那么这个值是比大顶堆所有元素都要大的。所以不会破坏堆之间数值大小的条件, 然后将小顶堆堆顶元素(这个元素大于大顶堆所有值, 小于小顶堆其他所有值), 所以将其直接插入到大顶堆当中, 这个元素就成为大顶堆的最大元素。(2) 假如当前的元素, 小于所有小顶堆元素, 并且比某些大顶堆元素还要小, 先将其插入到小顶堆当中, 这个元素成为小顶堆的堆顶元素, 然后将此最小的堆顶元素, 再插入到大顶堆当中, 相当于直接插入到大顶堆当中, 效果是一样的。

如果要插入到小顶堆当中, 也是同样的道理: 先插入到大顶堆, 然后将大顶堆最大元素插入到小顶堆。

// 最大堆 用于存储左半边数据
private Queue<Integer> maxHeap;
// 最小堆 用于存储右半边数据
private Queue<Integer> minHeap;
// 标记总大小, 用于判断当前元素是要插入到大顶堆当中还是小顶堆当中
private int size;

/** initialize your data structure here. */
public MedianFinder() {
    maxHeap = new PriorityQueue<>((x, y) -> (y - x));
    minHeap = new PriorityQueue<>();
}

public void addNum(int num) {
    if ((size & 0x1) == 0) {
        // 总数为偶数, 需要往最大堆插数据
        minHeap.add(num);
        maxHeap.add(minHeap.poll());
    } else {
        // 总数为奇数 需要往最小堆插数据
        maxHeap.add(num);
        minHeap.add(maxHeap.poll());
    }
    size++;
}

public double findMedian() {
    if ((size & 0x1) == 1) {
        // 总数为奇数时, 中位数为最大堆堆顶
        return maxHeap.peek();
    } else {
        // 总数为偶数时, 中位数为两个堆的堆顶元素的平均值
        return (maxHeap.peek() + minHeap.peek()) * 1.0 / 2;
    }
}