剑指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、findMedia进行 50000 次调用。

题解

可以维护两个堆，左边为最大堆，右边为最小堆。最大堆的堆顶元素不能大于最小堆的堆顶元素。

class MedianFinder {
    PriorityQueue<Integer> maxHeap, minHeap;

    /** initialize your data structure here. */
    public MedianFinder() {
        this.maxHeap = new PriorityQueue<>((o1, o2) -> o2 - o1);
        this.minHeap = new PriorityQueue<>();
    }
    
    public void addNum(int num) {
        // 添加数后要保证最小堆和最大堆的size差不能超过1
        // 而且最小堆中的堆顶元素要比最大堆中的堆顶元素要大
        if (maxHeap.size() == minHeap.size()) {
            maxHeap.offer(num);
            // 平衡操作
            minHeap.offer(maxHeap.poll());
        } else {
            // 不等的情况肯定是最小堆size大于最大堆
            minHeap.offer(num);
            // 平衡操作
            maxHeap.offer(minHeap.poll());
        }
    }
    
    public double findMedian() {
        if (maxHeap.size() == 0 && minHeap.size() == 0) {
            return 0.0;
        }
        // 如果两个堆size相等，返回中位数。如果不等，那么肯定是最小堆比最大堆size大，返回最小堆顶即可
        return maxHeap.size() == minHeap.size() ? (maxHeap.peek() + minHeap.peek()) * 0.5 : (double)minHeap.peek();
    }
}

取中位数时间复杂度：$O(1)$。
直接取堆顶元素即可。

存数据时间复杂度：$O(log(N))$。
$N$为当前元素个数。存数据涉及到3次出堆入堆操作，每次耗时$O(log(N))$。