没有SortedList，如何更快找到中间值

做了这么多道题，我们发现总是需要某种数据结构辅助我们解决一些问题，算法跟数据结构总是不分家。一般我们使用的语言都会给我们内置常用的数据结构，堆啊栈啊列表啊等等，用多了的人对于它们的作用想必还是比较清楚的。

我最前两天刷题遇到这样一个题目：设计一个类去计算一个数字流的中值。

这个类要包含下面两个方法：

1. insertNum(int num)数字流中新增一个数字。
2. findMedian()返回当前被增加的数字们的中值，如果数字个数是偶数，返回中间两个数的平均值。

这道题目乍一看很简单，简单中透露着一丝危险的味道。首先我想到的是把所有元素存进一个SortedList里，然后找中值也不是很难的事情。但是想在SortedList中插入一个元素时间复杂度是O(N)，作为一个Hard的题目它不该如此简单，有木有什么办法可以做得更好？仔细想来我们只是想获得中值或者说最中间的两个数，并不是要给所有数字排序，其他的数字我们不太关心，但是不排序要怎么找到中间的数字呢？

先来分析一波，假设X是一个列表的中值，这意味着这个列表中一般的数字小于等于X，另一半大于等于X。这让我们把一个列表分成了两半，一半存储比它小的数字（暂且叫它smallNumList），一半存储比它大的数字（暂且叫它largeNumList），那整个列表的中值就在smallNumList的最大值或者largeNumList的最小值产生，对不对？

要找最大或者最小，那第一个跳入脑海中的数据结构是堆。堆很多人都知道的，可以帮助我们快速找到最大或是最小的元素。今天我们的场景还比较特殊，它既要最大，也要最小，它需要两个堆才能完成。

前面的思路已经决定了主要方向，现在就是要来看看我们具体怎么利用堆了：

1. 我们可以把第一半数字放进一个Max Heap（也就是smallNumList），因为我们想要知道第一部分的最大值。
2. 我们可以把第二部分放进Min Heap（也就是largeNumList），这儿我们需要找到一个最小值。
3. 向堆中插入一个元素的时间复杂度是O(logN)，是比我们直接使用SortedList要快的。
4. 任何时候我们需要计算中值时只要要用两个堆最顶部的元素来判断就好了。
5. 我们要要让两个堆的元素数量保持平衡，一半一半，这样才能正确找到中值，如果数字的数量是奇数，我们就把它放在MaxHeap里面，这时候中值就是它的顶部元素。

这波分析下来，答案已经呼之欲出了：

class MedianOfAStream {

    PriorityQueue<Integer> maxHeap; //包含数字的前半段
    PriorityQueue<Integer> minHeap; //包含数字的后半段

    public MedianOfAStream() {
        maxHeap = new PriorityQueue<>((a, b) -> b - a);
        minHeap = new PriorityQueue<>((a, b) -> a - b);
    }

    public void insertNum(int num) {
        if (maxHeap.isEmpty() || maxHeap.peek() >= num)
            maxHeap.add(num);
        else
            minHeap.add(num);

        // 平衡两个堆的数量
        if (maxHeap.size() > minHeap.size() + 1)
            minHeap.add(maxHeap.poll());
        else if (maxHeap.size() < minHeap.size())
            maxHeap.add(minHeap.poll());
    }

    public double findMedian() {
        if (maxHeap.size() == minHeap.size()) {
            // 偶数取平均值
            return maxHeap.peek() / 2.0 + minHeap.peek() / 2.0;
        }
        // 奇数取 max-heap 最大值
        return maxHeap.peek();
    }
}

一套行云流水的操作下来，我们的insertNum方法时间复杂度已经到了O(logN)，findMedian时间复杂度只有O(1)了。不过空间复杂度没有变还是O(N)，我们还是需要保存每个数字。

趁热打铁我又赶紧来了一道相关的题：给定一个数字数组跟一个数字k，找出这个数组所有大小为k的字数组的中值。

这个大小为k的子数组的限制让我想起来之前讨论过的滑动窗口，而滑动窗口也类似于一个动态的数字流，进一个出一个，这大概是用双堆没跑了。真有意思，能把我们学的几种套路结合起来考察，巩固新的同时复习旧的，美滋滋。

乐呵完了我们赶紧来分析一下这道题的注意点。我们需要维护一个大小为k的滑动窗口，每次插入一个元素就要有一个元素出去，移除一个元素之后，我们还要注意维持两边的数量。其它好像没什么要注意的地方，还按之前的逻辑来应该走得通，我们试试。

class SlidingWindowMedian {
    PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
    PriorityQueue<Integer> minHeap = new PriorityQueue<>();

    public double[] findSlidingWindowMedian(int[] nums, int k) {
        double[] result = new double[nums.length - k + 1];
        for (int i = 0; i < nums.length; i++) {
            if (maxHeap.size() == 0 || maxHeap.peek() >= nums[i]) {
                maxHeap.add(nums[i]);
            } else {
                minHeap.add(nums[i]);
            }
            rebalanceHeaps();

            if (i - k + 1 >= 0) { // 判断是否达到窗口k 的数量
                // 把中值加入结果列表
                if (maxHeap.size() == minHeap.size()) {
                    // 这里要计算平均数
                    result[i - k + 1] = maxHeap.peek() / 2.0 + minHeap.peek() / 2.0;
                } else { // 这里直接拿maxHeap的顶部元素
                    result[i - k + 1] = maxHeap.peek();
                }

                // 滑动窗口
                int elementToBeRemoved = nums[i - k + 1];
                if (elementToBeRemoved <= maxHeap.peek()) {
                    maxHeap.remove(elementToBeRemoved);
                } else {
                    minHeap.remove(elementToBeRemoved);
                }
                rebalanceHeaps();
            }
        }
        return result;
    }

    private void rebalanceHeaps() {
        // max-heap最多比min-heap多一个元素
        if (maxHeap.size() > minHeap.size() + 1)
            minHeap.add(maxHeap.poll());
        else if (maxHeap.size() < minHeap.size())
            maxHeap.add(minHeap.poll());
    }

}

这个套路有个显而易见的名字叫双堆法。这边只要注意一下滑动窗口的注意事项跟保持两边的数目一致，其它没什么大问题，都是我们上面讨论的基本思想，兄弟们，总结出题型的基本思想很重要哇。

总的看下来这类题其实也不难，主要在于你知不知道有这样一种数据结构，你知不知道它可以让你更快。大家在刷题前可以先花点时间把常用的数据结构再复习一下，你都不知道有这个东西，那想破脑袋也想不到可以用它去解决问题。工欲善其事必先利其器，数据结构就是我们的工具，一定要先花时间好好了解我们的工具。Happy coding~