LeetCode-480 Sliding Window Median

Description

给出数组\(nums[]\)，窗口大小\(k\)，求滑动窗口中的（排序后）中值。具体见下面样例：

Window position                Median
---------------                -----
[1  3  -1] -3  5  3  6  7        1
 1 [3  -1  -3] 5  3  6  7       -1
 1  3 [-1  -3  5] 3  6  7       -1
 1  3  -1 [-3  5  3] 6  7        3
 1  3  -1  -3 [5  3  6] 7        5
 1  3  -1  -3  5 [3  6  7]       6

Constraints

\(nums[]\)长度不超\(10^5\)
\(nums[i]\)数值在整型int范围内

Solution

两个对顶堆的事情，麻烦在于corner case。这里可以用不变式来避免多种情况的讨论，设不变式为:

\( size(small\ set) \ge size(large\ set) \)
\( size(small\ set) \le size(large\ set) + 1 \)

也就是一般2个集合都是等同大小的，最多允许\(small\ set\)多出1个元素

Code

class Solution {
public:
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        // small set & large set
        auto ss = multiset<int, less<int>> {};
        auto ls = multiset<int, greater<int>> {};

        auto pop = [&](auto &q) {
            auto iter = --q.end();
            auto val = *iter;
            q.erase(iter);
            return val;
        };

        vector<double> ans;

        for(int i = 0; i < nums.size(); ++i) {
            auto v = nums[i];
            ss.emplace(v);

            // 不管怎样，先交换一下
            if(!ls.empty()) {
                auto x = pop(ss);
                auto y = pop(ls);
                ss.emplace(min(x, y));
                ls.emplace(max(x, y));
            }

            // 如果存在违反不变式
            if(ss.size() > ls.size() + 1) {
                auto e = pop(ss);
                ls.emplace(e);
            }

            if(i >= k-1) {
                if(ss.size()+ls.size() & 1) {
                    ans.emplace_back(*--ss.end());
                } else {
                    ans.emplace_back((0.0 + *--ss.end() + *--ls.end()) / 2);
                }

                auto evict = nums[i - (k - 1)];
                // 考虑到下一次还是从ss里插入的，这样不会违反不变式
                if(auto siter = ss.find(evict); siter != ss.end()) {
                    ss.erase(siter);
                // 如果实在没办法，那就需要往ss里挪用一个
                // 下一次插入ss后仍然满足不变式
                } else {
                    auto liter = ls.find(evict);
                    ls.erase(liter);
                    ls.emplace(pop(ss));
                }
            }
        }

        return ans;
    }
};

Link

LeetCode-480 Sliding Window Median