想要精通算法和SQL的成长之路 - 滑动窗口和大小根堆

前言

想要精通算法和SQL的成长之路 - 系列导航

一. 大小根堆

先来说下大小根堆是什么：

• 大根堆：栈顶元素最大（上图左侧部分），栈底至栈顶元素值递增。
• 小根堆：栈顶元素最小（上图右侧部分），栈底至栈顶元素值递减。

在Java当中，可以用什么来表示大小根堆？

小根堆：

Queue<Integer> small = new PriorityQueue<>();
// 或者 x - y 是计算，在特殊情况下可能造成精度越界的情况
Queue<Integer> small = new PriorityQueue<>((x, y) -> x - y);
// 或者，Integer.compare 是纯比较，不会出现精度越界
Queue<Integer> small = new PriorityQueue<>((x, y) -> Integer.compare(x, y));
// 或者
Queue<Integer> small = new PriorityQueue<>(Integer::compare);

大根堆：

Queue<Integer> big = new PriorityQueue<>((x, y) -> y - x);

大小根堆的常规操作：

• 获取栈顶元素：peek();
• 栈顶元素移除：poll();

二. 数据流的中位数

原题链接

再说下我们的思路：

1. 同时维护大小根堆，并且约定小根堆的元素个数总是 >= 大根堆元素个数（最多个数多一个）。
2. 如果元素个数是奇数，那么中位数就是小根堆堆顶元素。
3. 如果元素个数是偶数，那么中位数就是（大根堆堆顶 + 小根堆堆顶） / 2。

1.1 初始化

Queue<Integer> big, small;/*** big                      small* 最小值 ---> 大根堆顶 中位数 小根堆顶 ---> 最大值*/
public MedianFinder() {small = new PriorityQueue<>();// 小根堆，堆顶元素最小（存储比中位数大的部分）big = new PriorityQueue<>((x, y) -> y - x);// 大根堆，堆顶元素最大（存储比中位数小的部分）
}

1.2 插入操作

插入的时候，我们考虑到两种情况：

• 如果大小根堆的元素个数相等，我们优先把新元素加入到小根堆。
• 否则，将元素加入到大根堆。

但是，我们并不知道以下三者的关系：

• 大根堆堆顶元素值。
• 当前待加入元素值。
• 小根堆堆顶元素值。

而我们需要去维护他们，一定满足：大根堆堆顶元素值 < 小根堆堆顶元素值。

咋办呢？以第一种情况为例，我们可以：

• 先把元素加入到大根堆。那么经过排序后，大根堆的堆顶元素就是最大的那个（可能是当前元素，也可能不是）。此时大根堆Size > 小根堆Size。
• 把大根堆堆顶元素移除，加入到小根堆。小根堆经过排序后，这样就能保证大根堆堆顶元素值 < 小根堆堆顶元素值。

写成代码就是：

public void addNum(int num) {// 如果大小根堆 的 大小 一样，我们往小根堆放元素。让小根堆size >= 大根堆sizeif (big.size() == small.size()) {// 方式一定是先让放大根堆，再把大根堆的堆顶元素移除到小根堆big.add(num);small.add(big.poll());} else {small.add(num);big.add(small.poll());}
}

1.3 完整代码

那么查询函数就更简单了，结合上面的思路，我们得到完整代码如下：

public class MedianFinder {Queue<Integer> big, small;/*** big                      small* 最小值 ---> 大根堆顶 中位数 小根堆顶 ---> 最大值*/public MedianFinder() {small = new PriorityQueue<>();// 小根堆，堆顶元素最小（存储比中位数大的部分）big = new PriorityQueue<>((x, y) -> y - x);// 大根堆，堆顶元素最大（存储比中位数小的部分）}public void addNum(int num) {// 如果大小根堆 的 大小 一样，我们往小根堆放元素。让小根堆size >= 大根堆sizeif (big.size() == small.size()) {// 方式一定是先让放大根堆，再把大根堆的堆顶元素移除到小根堆big.add(num);small.add(big.poll());} else {small.add(num);big.add(small.poll());}}public double findMedian() {return small.size() == big.size() ? (small.peek() + big.peek()) / 2.0 : small.peek();}
}

三. 滑动窗口中位数

原题链接

思路如下：

1. 我们先创建一个窗口，把前k个数字通过大小根堆的方式去维护（题目一的思路）。
2. 后续每次滑动窗口的移动，都带来两个变数：一个旧元素会从窗口出移除（但是从大根堆移除还是小根堆移除？），一个新元素会加入到窗口中（加入到大根堆还是小根堆？）
3. 由于第二步的变数，可能导致大小根堆的Size不均衡。我们的目的：让小根堆的Size >= 大根堆Size，最多多一个元素。
4. 因此每次滑动窗口的移动，我们还需要维护大小根堆。

3.1 在第一题的基础上改造

首先考虑到精度的问题，我们的大小根堆不能在根据差值来比较了，而是：

right = new PriorityQueue<>((x, y) -> Integer.compare(x, y));// 小根堆，堆顶元素最小（存储比中位数大的部分）
left = new PriorityQueue<>((x, y) -> Integer.compare(y, x));// 大根堆，堆顶元素最大（存储比中位数小的部分）

其次，求中位数的时候，也需要大小根堆的堆顶元素，先除以2，再和相加：

if (left.size() == right.size()) {return (left.peek() / 2.0) + (right.peek() / 2.0);

最终代码如下：

public class Test480 {Queue<Integer> left, right;public double[] medianSlidingWindow(int[] nums, int k) {right = new PriorityQueue<>((x, y) -> Integer.compare(x, y));// 小根堆，堆顶元素最小（存储比中位数大的部分）left = new PriorityQueue<>((x, y) -> Integer.compare(y, x));// 大根堆，堆顶元素最大（存储比中位数小的部分）int len = nums.length;// 结果集double[] res = new double[len - k + 1];// 创建大小根堆for (int i = 0; i < k; i++) {right.add(nums[i]);}for (int i = 0; i < k / 2; i++) {left.add(right.poll());}// 初始化第一个中位数res[0] = findMedian();for (int i = k; i < len; i++) {// 滑动窗口长度固定，每次移动，都有一个元素要删除和一个元素要新加入int del = nums[i - k], add = nums[i];if (add >= right.peek()) {right.add(add);} else {left.add(add);}// 如果待删除元素在小根堆,在小根堆处删除，否则在大根堆中删除if (del >= right.peek()) {right.remove(del);} else {left.remove(del);}// 维护大小根堆的元素个数adjust();res[i - k + 1] = findMedian();}return res;}void adjust() {while (left.size() > right.size()) {right.add(left.poll());}while (right.size() - left.size() > 1) {left.add(right.poll());}}public double findMedian() {if (left.size() == right.size()) {return (left.peek() / 2.0) + (right.peek() / 2.0);} else {return right.peek() * 1.0;}}
}

这个写法其实是没问题的，但是在元素个数非常大的情况下，就容易超时：

3.2 栈的remove操作

问题处在优先队列的的一个元素remove操作：

它是先查找(复杂度O(N))，再进行删除(复杂度O(logN))，所以会超时。因此我们这里可以引入红黑树来进行替代。

有这么几个需要注意的地方：

1. 我们用TreeSet存储元素的时候，不再是元素值，而是元素的下标。 因为题目中同一个窗口的元素可能重复。元素值相等的时候，根据下标大小来比较。

Comparator<Integer> comparator = (x, y) -> nums[x] != nums[y] ? Integer.compare(nums[x], nums[y]) : x - y;
right = new TreeSet<>(comparator);// 小根堆，堆顶元素最小（存储比中位数大的部分）
left = new TreeSet<>(comparator.reversed());// 大根堆，堆顶元素最大（存储比中位数小的部分）

2. 滑动窗口移动的时候。需要删除对应的元素下标 ，由于存在重复值，我们需要大小根堆都把这个下标给剔除。
3. 原peek函数替代为first函数。poll函数替代为pollFirst函数。

完整代码如下：

public class Test480 {TreeSet<Integer> left, right;int[] nums;public double[] medianSlidingWindow(int[] nums, int k) {this.nums = nums;Comparator<Integer> comparator = (x, y) -> nums[x] != nums[y] ? Integer.compare(nums[x], nums[y]) : x - y;right = new TreeSet<>(comparator);// 小根堆，堆顶元素最小（存储比中位数大的部分）left = new TreeSet<>(comparator.reversed());// 大根堆，堆顶元素最大（存储比中位数小的部分）int len = nums.length;// 结果集double[] res = new double[len - k + 1];// 创建大小根堆for (int i = 0; i < k; i++) {addToWindow(i);}res[0] = findMedian();for (int i = k; i < len; i++) {// 滑动窗口长度固定，每次移动，都有一个元素要删除和一个元素要新加入left.remove(i - k);right.remove(i - k);addToWindow(i);res[i - k + 1] = findMedian();}return res;}void addToWindow(int index) {// 我们总是把新元素先统一加入到大根堆。right.add(index);left.add(right.pollFirst());// 然后再维护大小while (left.size() > right.size()) {right.add(left.pollFirst());}}public double findMedian() {if (left.size() == right.size()) {return (nums[left.first()] / 2.0) + (nums[right.first()] / 2.0);} else {return nums[right.first()] * 1.0;}}
}