数组中滑动窗口的中位数 | 套装2

题目描述

给定一个大小为 n 的数组，有一个大小为 k 的滑动窗口从数组的最左边移动到最右边。从左到右移动窗口，每次窗口向右移动一个位置。例如，如果数组为 [1,3,-1,-3,5,3,6,7]，并且 k = 3，则所有滑动窗口的中位数为 [1,-1,-1,3,5,6]。

解题思路

方法一：暴力法

对于每个滑动窗口，我们可以直接将其所有元素排序，然后直接找出中位数。时间复杂度为 $O(n k \log{k})$。

方法二：堆

对于这个问题，可以用一个数据结构来维护滑动窗口中的元素，并且能够在 $O(\log{k})$ 时间内找到中位数。其中，本套装提供两种不同的堆的实现方式。

方法二.1：大根堆 + 小根堆

我们可以使用两个优先队列，一个是小根堆，一个是大根堆。小根堆维护较大的一半数，大根堆维护较小的一半数。当元素个数为偶数时，我们将这个新数插入到小根堆。如果小根堆的元素个数变成多于大根堆的个数 (也就是目前的中位数在小根堆中)，我们从小根堆中弹出堆顶元素并插入到大根堆中；否则插入到大根堆中。当元素个数为奇数时，我们将这个数插入到大根堆。如果大根堆的元素个数多于小根堆的个数，则从大根堆中弹出堆顶元素插入到小根堆中；否则插入到小根堆中。这样保证了小根堆的元素个数始终大于等于大根堆的元素个数。中位数就可以从堆顶元素中取得。

方法二.2：Multiset

Multiset 是一个可以重复的有序集合。对于这个问题，我们可以使用两个 multiset 来维护滑动窗口中的元素，一个是前 $\lfloor \frac{k}{2} \rfloor$ 小的有序集合，一个是后 $\lceil \frac{k}{2} \rceil$ 大的有序集合。

代码实现

方法一：暴力法

class Solution {
public:
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        vector<double> res;
        for (int i = 0; i <= nums.size() - k; ++i) {
            vector<int> window(nums.begin() + i, nums.begin() + i + k);
            sort(window.begin(), window.end());
            if (k % 2 == 1) {
                res.push_back(window[k / 2]);
            } else {
                res.push_back((double)(window[k / 2 - 1] + window[k / 2]) / 2);
            }
        }
        return res;
    }
};

方法二.1：大根堆 + 小根堆

class Solution {
public:
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        multiset<int, greater<int>> small;
        multiset<int> large;
        vector<double> res;
        for (int i = 0; i < nums.size(); ++i) {
            if (i >= k) {
                if (small.count(nums[i - k])) {
                    small.erase(small.find(nums[i - k]));
                } else {
                    large.erase(large.find(nums[i - k]));
                }
            }
            if (small.size() <= large.size()) {
                if (large.empty() || nums[i] <= *large.begin()) {
                    small.insert(nums[i]);
                } else {
                    small.insert(*large.begin());
                    large.erase(large.begin());
                    large.insert(nums[i]);
                }
            } else {
                if (nums[i] >= *small.begin()) {
                    large.insert(nums[i]);
                } else {
                    large.insert(*small.begin());
                    small.erase(small.begin());
                    small.insert(nums[i]);
                }
            }
            if (i >= k - 1) {
                if (k % 2 == 1) {
                    res.push_back(*small.begin());
                } else {
                    res.push_back((double)(*small.begin() + *large.begin()) / 2);
                }
            }
        }
        return res;
    }
};

方法二.2：Multiset

class Solution {
public:
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        multiset<int> small(nums.begin(), nums.begin() + k / 2);
        multiset<int> large(nums.begin() + k / 2, nums.begin() + k);
        vector<double> res;
        for (int i = k; i <= nums.size(); ++i) {
            if (k % 2 == 1) {
                res.push_back(*small.rbegin());
            } else {
                res.push_back((double)(*small.rbegin() + *large.begin()) / 2);
            }
            if (i == nums.size()) {
                break;
            }
            if (small.count(nums[i - k])) {
                small.erase(small.find(nums[i - k]));
                small.insert(*large.begin());
                large.erase(large.begin());
            } else {
                large.erase(large.find(nums[i - k]));
                large.insert(*small.rbegin());
                small.erase(--small.end());
            }
            if (small.empty() || nums[i] <= *small.rbegin()) {
                small.insert(nums[i]);
            } else {
                large.insert(nums[i]);
            }
        }
        return res;
    }
};

总结

这道题使用堆的方法是经典的滑动窗口中位数求解方法，了解用堆解决滑动窗口问题的同学需要掌握。此外，多种解法的实现也提供了多样化的思路，希望大家在实践中探究，丰富解题思维。