如何求解中位数

如果输入一个数组，让你求中位数，这个比较好办，排序，如果数组长度是奇数，最中间的一个元素就是中位数，如果是偶数，最中间两个元素的平均数作为中位数。

如果数据规模非常巨大，排序是不太现实的，那么可以使用概率算法，随机抽取一部分数据，排序，求中位数，作为所有数据的中位数。

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
进阶:
如果数据流中所有整数都在 0 到 100 范围内，你将如何优化你的算法？
如果数据流中 99% 的整数都在 0 到 100 范围内，你将如何优化你的算法？
来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/find-median-from-data-stream
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

一个直接的解法就是用一个数组记录所有 addNum 添加进来的数字，通过插入排序的逻辑保证数组中的元素有序，当调用 findMedian 方法时，可以通过数组索引直接计算中位数。

但是用数组作为底层容器的问题也很明显，addNum 搜索插入位置的时候可以使用二分搜索算法，但是插入操作需要搬动数据，所以最坏的时间复杂度为 O(N)。

解题思路

需要有序数组结构，本题的核心思路是使用两个优先队列。

中位数是有序数组最中间的元素算出来的。可以将有序数组抽象成一个倒三角形，宽度可以视为元素的大小，那么这个倒三角形的中部就是计算中位数的元素：如图所示

然后把这个大的倒三角形从正中间切成两半，变成一个小倒三角形和一个梯形，这个小倒三角形相当于从小到大的有序数组，梯形相当于一个从大到小的有序数组。

中位数就可以通过小倒三角形和梯形顶部的元素算出来。其中，小倒三角形就是一个大顶堆，梯形就是小顶堆，中位数可以通过他们的堆顶元素算出来。

当然，这两个堆需要算法逻辑正确维护，才能保证堆顶元素是可以算出正确的中位数。其中，很容易看出，两个堆中的元素之差不能超过 1。

因为要求的是中位数，假设元素总数是 n，如果 n 是偶数，我们希望两个堆的元素个数相同，这样把两个堆的堆顶元素拿出来求个平均数就是中位数；如果 n 是奇数，我们希望两个堆的元素个数分别是 n / 2 和 n / 2。这样元素多的那个堆的堆顶元素就是中位数。

根据这个逻辑，可以直接写出 findMedian 函数的代码：

class MedianFinder {
private:
    // 小顶堆: 升序
    priority_queue<int, vector<int>, greater<int>> large;
    // 大顶堆：降序
    priority_queue<int, vector<int>, less<int>> small;
public:
    MedianFinder() {

    }

    double findMedian() {
        // 如果元素不一样多，多的那个堆的堆顶元素就是中位数
        if (large.size() < small.size()) {
            return small.top();
        } else if (large.size() > small.size()) {
            return large.top();
        }
        return (double)(large.top() + small.top()) / 2;
    }
};

现在的问题是，如何实现 addNum 方法，维护两个堆中的元素之差不能超过 1 这个条件？

一个简单的想法是：每次调用 addNum 函数的时候，比较一下 large 和 small 的元素个数，谁的元素少我们就加到谁那里，如果它们的元素一样多，默认加到 large 里：

void addNum(int num) {
    if (large.size() <= small.size()) {
        large.push(num);
    } else {
        small.push(num);
    }
}

如果依次加入元素 1, 2, 3。则两个堆中的元素分布如下图：

从上图得知，不仅要维护 large 和 small 的元素之差不超过 1，还要维护 large 堆顶元素要大于等于 small 堆顶元素。逻辑上：想要往 large 里添加元素，不能直接添加，而是要先往 small 里添加，然后再把 small 的堆顶元素加到 large 中；向 small 中添加元素同理。代码如下：

void addNum(int num) {
    if (large.size() <= small.size()) {
        small.push(num);
        large.push(small.top());
        small.pop();
    } else {
        large.push(num);
        small.push(large.top());
        small.pop();
    }
}

比如，准备向 large 中插入元素：

如果插入的 num 小于 small 的堆顶元素，那么 num 就会留在 small 堆中，为了保证两个堆中的元素数量之差不大于 1，作为交换，把 small 堆顶的元素再插入到 large 堆里。

如果插入的num大于small的堆顶元素，那么num就会成为samll的堆顶元素，最后还是会被插入large堆中