[leetcode题解] 第480题Sliding Window Median

https://leetcode-cn.com/problems/sliding-window-median/

分析

这道题是295.数据流的中位数这题的进阶版，此时数据流变成了一个固定大小的窗口，每次新加入一个数字就会有一个旧的数字被移除，因此，这相当于在295题的基础上引入了删除操作。
295题的解法是用两个堆来存放数据流中的数据集，small存小的那一半，large存大的那一半，并保证smallSize == largeSize || smallSize == largeSize + 1。这样一来获取中位数就是$O(1)$的操作，简单不需要赘述；新加入一个数字时需要维护这两个堆满足的条件，有两种写法：

if (left.size() == right.size()) {
    right.push(num);
    left.push(right.top());
    right.pop();
}
else {
    left.push(num);
    right.push(left.top());
    left.pop();
}

if (small.empty() || num <= small.top()) small.push(num);
else large.push(num);
// 堆内的元素个数不满足条件的两种情况
if (small.size() == large.size() + 2) {
	large.push(small.top());
    small.pop();
}
else if (small.size() + 1 == large.size()) {
	small.push(large.top());
	large.pop();
}

显然第一种写法更简洁，这里提到第二种写法是因为它适合用在本题的编码。

解法

相比于295题，本题多了一个删除操作，而堆是只能删除堆顶元素的，因此看上去不可实现。为了解决此题，我们可以引入一个延迟删除的概念，即当我们需要把一个数字移出窗口时，不直接把它从堆内删除（也无法实现），而是用哈希表记录这个删除操作，当这个数字成为堆顶时，才真正执行删除操作。
引入一个辅助函数prune来帮助执行删除操作：

template <typename T>
void prune(T& queue) {
    while (!queue.empty() && delayed[queue.top()] > 0) {
        delayed[queue.top()]--;
        queue.pop();
    }
}

为了容易实现编码，我们需要维护这样一个状态：addNum、earseNum以及getMedian这3种操作在执行之前、执行之后，都要保证两个堆顶的元素是有效的（无需延迟删除）。
getMedian操作显然不会影响堆顶元素的有效性，因此不用修改。

addNum

为了维护状态，addNum需要在哪些地方调用prune呢？首先把addNum改写如下：

if (small.empty() || num <= small.top()) small.push(num), smallSize++;
else large.push(num), largeSize++;
if (smallSize == largeSize + 2) {
     large.push(small.top());
     small.pop();
     largeSize++;
     smallSize--;
 }
 else if (smallSize + 1 == largeSize) {
     small.push(large.top());
     large.pop();
     largeSize--;
     smallSize++;
 }

相比前面那种写法，这里引入了smallSize和largeSize两个变量，这是自然的，因为small.size()不再代表small实际包含的元素个数了，这里面还有需要被延迟删除的元素。
下面逐行分析哪些代码会导致堆顶元素变得无效：

if (small.empty() || num <= small.top()) small.push(num), smallSize++;

执行此行代码，由于num <= small.top()而且small是大根堆，那么small的堆顶元素不变，依然有效。

else large.push(num), largeSize++;

此时small.top() < num <= large.top()，当num == large.top()时堆顶不变，依然有效；当small.top() < num < large.top()时，易知small和large里面之前都是没有num的，自然就不会需要被延迟删除，因此num会是一个有效的堆顶元素。

large.push(small.top());

这行代码会使small.top()成为large的堆顶元素，由于small.top()一定有效，那么large新的堆顶也是有效的。

small.pop();

如果small中第二大的元素是要被延迟删除的，那么这行代码会导致small的堆顶变得无效。
综上，仅需要在small.pop()后面加上prune(small)。

eraseNum

eraseNum是新引入的操作，当要删除的元素正好是堆顶元素时，就需要立即删除，否则延迟删除（仅更新哈希表）：

delayed[num]++;
if (num <= small.top()) {
    --smallSize;
    if (num == small.top()) prune(small);
}
else {
    --largeSize;
    if (num == large.top()) prune(large);
}
if (smallSize == largeSize + 2) {
    large.push(small.top());
    small.pop();
    largeSize++;
    smallSize--;
}
else if (smallSize + 1 == largeSize) {
    small.push(large.top());
    large.pop();
    largeSize--;
    smallSize++;
}

易知需要在哪些地方加上prune函数。第二个if/else是维护两个堆内元素数量平衡的，在addNum中也有用到，因此代码可以复用。

代码

class Solution {
private:
    priority_queue<int> small;
    priority_queue<int, vector<int>, greater<int>> large;
    unordered_map<int, int> delayed;

    //small和large实际包含的元素(去除需要延迟删除的元素)
    int smallSize = 0, largeSize = 0;

    void addNum(int num) {
        if (small.empty() || num <= small.top()) small.push(num), smallSize++;
        else large.push(num), largeSize++;
        makeBalance();
    }

    void eraseNum(int num) {
        delayed[num]++;
        if (num <= small.top()) {
            --smallSize;
            if (num == small.top()) prune(small);
        }
        else {
            --largeSize;
            if (num == large.top()) prune(large);
        }
        makeBalance();
    }

    void makeBalance() {
        if (smallSize == largeSize + 2) {
            large.push(small.top());
            small.pop();
            largeSize++;
            smallSize--;
            prune(small);
        }
        else if (smallSize + 1 == largeSize) {
            small.push(large.top());
            large.pop();
            largeSize--;
            smallSize++;
            prune(large);
        }
    }

    template <typename T>
    void prune(T& queue) {
        while (!queue.empty() && delayed[queue.top()] > 0) {
            delayed[queue.top()]--;
            queue.pop();
        }
    }

    double getMedian(int k) {
        return k & 1 ? small.top() : ((double)small.top() + large.top()) / 2.0;
    }
public:
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        int n = nums.size();
        if (!n) return {};
        vector<double> ans;
        for (int i = 0; i < k; ++i) addNum(nums[i]);
        ans.emplace_back(getMedian(k));
        for (int i = k; i < n; ++i) {
            addNum(nums[i]);
            eraseNum(nums[i - k]);
            ans.emplace_back(getMedian(k));
        }
        return ans;
    }
};