题目描述

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

例如，

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

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

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

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

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/find-median-from-data-stream\

样例

样例

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

算法与数据结构

堆

题解

一种显然的做法是，将数字存储 ArrayList 中，每次需要输出中位数时，对其进行排序，并输出中间值。

对于一组静态数据，中位数是固定的，我们可以先排序，第 $\frac{n}{2}$ 个数据就是中位数。不论询问多少次中位数，直接返回这个固定的值就好了。所以，尽管排序的代价比较大，但是边际成本会很小。

但是，如果我们面对的是动态数据集合，中位数在不停地变动，如果再用先排序的方法，每次询问中位数的时候，都要先进行排序，那效率就不高了。

一种显然的做法是使用优先级队列，对于一个优先级队列，每次遇到 addNum() 就往里面加入新的数据，优先级队列会自动帮我们从小到大排序，而每次遇到 findMedian() 的时候，就输出 pq.size() / 2 位置上的值（如果 size 为奇数）或者 pq.size() / 2 与 pq.size() / 2 + 1 这两个位置的平均值（如果 size 为偶数）。

然而，对于一般的优先级队列，只能访问队首的元素，并不提供随机访问的接口，这时候就只能手动维护两个堆了。

我们需要维护两个堆，一个大顶堆，一个小顶堆。大顶堆中存储前半部分数据，小顶堆中存储后半部分数据，且小顶堆中的数据都大于大顶堆中的数据。

也就是说，如果有 $n$ 个数据，$n$ 是偶数，我们从小到大排序，那前 $\frac{n}{2}$ 个数据存储在大顶堆中，后 $\frac{n}{2}$ 个数据存储在小顶堆中。如果 $n$ 是奇数，情况是类似的，大顶堆就存储 $\frac{n}{2} + 1$ 个数据，小顶堆中就存储 $\frac{n}{2}$ 个数据。

在寻找中位数的时候，如果 $n$ 是奇数，那么中位数就是大顶堆的堆顶元素，如果 $n$ 是偶数，那么中位数就是大顶堆和小顶堆各自堆顶元素的平均值。这个结论是显然成立的。

于是，findMedian() 就很容易得到了。

if (max.size() == min.size()) {
    // n 是偶数
    return 0.5 * (max.peek() + min.peek());
} else {
    // n 是奇数
    return max.peek();
}

这样，我们所有的工作就集中在了堆的调整上，如何在插入新数据的时候，依然保证上面的约定条件成立呢？

如果新加入的数据小于等于大顶堆的堆顶元素，我们就将这个新数据插入到大顶堆；否则，我们就将这个新数据插入到小顶堆。

if (num < max.peek()) {
    max.add(num);
} else {
    min.add(num);
}

这个时候就有可能出现两个堆中的数据个数不符合前面约定的情况：如果 $n$ 是偶数，两个堆中的数据个数都是 $\frac{n}{2}$；如果 $n$ 是奇数，大顶堆有 $\frac{n}{2} + 1$ 个数据，小顶堆有 $\frac{n}{2}$ 个数据。
这个时候，我们可以从一个堆中不停地将堆顶元素移动到另一个堆，通过这样的调整，来让两个堆中的数据满足上面的约定。

我们可以写一个循环，当大顶堆的元素个数小于小顶堆的元素个数时，将小顶堆的堆顶元素出堆，并加入大顶堆；反之，则将大顶堆的堆顶元素出堆，加入小顶堆。这个循环的终止条件是，两个堆的元素个数相等（$n$ 为偶数），或者大顶堆比小顶堆多 $1$ 个元素（$n$ 为奇数）。

while (true) {
    int minSize = min.size();
    int maxSize = max.size();
    if (minSize == maxSize || minSize + 1 == maxSize) {
        break;
    } else {
        if (minSize < maxSize) {
            min.add(max.poll());
        } else {
            max.add(min.poll());
        }
    }
}

最后检查一下 Corner Case。

对于本题来说，有可能出现错误的地方，就是当堆为空的时候，访问 peek() 的结果不能交给 int。

在 findMedian() 中，我们可以很自信地说出，这样的情况不可能存在。

因为既然是找中位数，则至少有 $1$ 个元素，而大顶堆的元素个数一定大于或者等于小顶堆的元素个数，因此，大顶堆元素个数不可能为 $0$，大顶堆不可能为空。

既然大顶堆不可能为空，那么，若 $n$ 为奇数，返回大顶堆的堆顶元素，就不可能遇到 null，而若 $n$ 为偶数，说明大顶堆和小顶堆的元素个数相等，那么小顶堆自然也不可能为空，因此 peek() 也是安全的，可以正常计算两个数的均值。

public double findMedian() {
    assert !max.isEmpty();
    if (max.size() == min.size()) {
        return 0.5 * (max.peek() + min.peek());
    } else {
        return max.peek();
    }
}

另外还有一个地方用到了 peek()，那就是在插入的地方。

上面说到，插入新的数据的时候，永远是与大顶堆的堆顶元素比较，而大顶堆不可能为空，所以也是安全的？

等等，那么，第一个元素怎么办？

简单做一下特判就可以了。

if (max.isEmpty() || num < max.peek()) {}

完整代码

class MedianFinder {

    private static PriorityQueue<Integer> min;
    private static PriorityQueue<Integer> max;

    public MedianFinder() {
        min = new PriorityQueue<>();
        max = new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2-o1;
            }
        });
    }

    public void addNum(int num) {
        if (max.isEmpty() || num < max.peek()) {
            max.add(num);
        } else {
            min.add(num);
        }
        while (true) {
            int minSize = min.size();
            int maxSize = max.size();
            if (minSize == maxSize || minSize + 1 == maxSize) {
                break;
            } else {
                if (minSize < maxSize) {
                    min.add(max.poll());
                } else {
                    max.add(min.poll());
                }
            }
        }
    }

    public double findMedian() {
        assert !max.isEmpty();
        if (max.size() == min.size()) {
            return 0.5 * (max.peek() + min.peek());
        } else {
            return max.peek();
        }
    }
}

备注

时间复杂度： $O(5 \cdot \log n) + O(1) \approx O(\log n)$。

最坏情况下，从顶部有三个堆插入和两个堆删除。每一个都需要花费 $O(\log n)$ 时间。

找到平均值需要持续的 $O(1)$ 时间，因为可以直接访问堆的顶部。

空间复杂度：$O(n)$ 用于在容器中保存输入的线性空间。

另外，在上文中提到了

if (pq.size() % 2 == 0) {
    return 0.5 * (pq.get(pq.size()) + pq.get(pq.size() + 1));
} else {
    return pq.get(pq.size() / 2);
}

自平衡二进制搜索树（如 AVL 树）具有一些非常有趣的特性。它们将树的高度保持在对数范围内。因此，插入新元素具有相当好的时间性能。中值总是缠绕在树根或它的一个子树上。大多数语言实现模拟这种行为的是 multiset 类。唯一的问题是跟踪中值元素。

我们保持两个指针：一个用于中位数较低的元素，另一个用于中位数较高的元素。当元素总数为奇数时，两个指针都指向同一个中值元素（因为在本例中只有一个中值）。当元素数为偶数时，指针指向两个连续的元素，其平均值是输入的代表中位数。