滑动窗口中位数 图解题解

这道题到底在问什么

输入

nums=[1,3,-1,-3,5,3,6,7], k=3

输出

[1,-1,-1,3,5,6]

最优解：一步一步想明白

1. 3记住这句:小的一半放大顶堆 small、大的一半放小顶堆 large,让 small 和 large 等大、或 small 多一个。k 奇数时中位数是 small 堆顶,k 偶数时取 small 顶和 large 顶的平均,都是 O(1) 读到。
2. 4先把第 0 个窗口建起来。第一个数 1 进 small(大顶堆),现在 small 堆顶是 1,large 还空着。
3. 5来第二个数 3。它比 small 堆顶 1 大,属于「较大的一半」,放进 large(小顶堆)。现在两边各一个,small 顶 1、large 顶 3。
4. 6第三个数 -1,它 < small 堆顶 1,属于「较小的一半」,进 small。这下 small 有 2 个、large 有 1 个,small 比 large 多一个,正好符合不变量,不用再平衡。
5. 7第 0 个窗口 [1, 3, -1] 就位:small 装较小的一半、堆顶是 1,large 装较大的一半。k 是奇数,中位数就是 small 堆顶 1。答案数组:[1]。
6. 8窗口右滑一格:离开窗口的是左边的 1,新进来的是右边的 -3。先把 1 从它所在的堆删掉,再把 -3 放进 small(它不大于 small 堆顶)。此刻 small 有 2 个、large 有 1 个。
7. 9检查两堆个数:small 2 个、large 1 个,已经符合 small 比 large 多 0 或 1 个,这一步不用搬动。
8. 10读 small 堆顶:-1,就是窗口 [3, -1, -3] 的中位数 -1。答案接成 [1, -1]。
9. 11窗口接着右移:离开窗口的是左边的 3,新进来的是右边的 5。先把 3 从它所在的堆删掉,再把 5 放进 large(它大于 small 堆顶)。此刻 small 有 2 个、large 有 1 个。
10. 12检查两堆个数:small 2 个、large 1 个,已经符合 small 比 large 多 0 或 1 个,这一步不用搬动。
11. 13读 small 堆顶:-1,就是窗口 [-1, -3, 5] 的中位数 -1。答案接成 [1, -1, -1]。
12. 14再滑一格:离开窗口的是左边的 -1,新进来的是右边的 3。先把 -1 从它所在的堆删掉,再把 3 放进 large(它大于 small 堆顶)。此刻 small 有 1 个、large 有 2 个。
13. 15两堆个数失衡了,做一次再平衡:large 多了一个,把 large 堆顶 3 挪到 small,直到 small 比 large 多 0 或 1 个。现在窗口 [-3, 5, 3] 的两堆排好了。
14. 16读 small 堆顶:3,就是窗口 [-3, 5, 3] 的中位数 3。答案接成 [1, -1, -1, 3]。
15. 17窗口继续右滑:离开窗口的是左边的 -3,新进来的是右边的 6。先把 -3 从它所在的堆删掉,再把 6 放进 large(它大于 small 堆顶)。此刻 small 有 1 个、large 有 2 个。
16. 18两堆个数失衡了,做一次再平衡:large 多了一个,把 large 堆顶 5 挪到 small,直到 small 比 large 多 0 或 1 个。现在窗口 [5, 3, 6] 的两堆排好了。
17. 19读 small 堆顶:5,就是窗口 [5, 3, 6] 的中位数 5。答案接成 [1, -1, -1, 3, 5]。
18. 20最后滑一格:离开窗口的是左边的 5,新进来的是右边的 7。先把 5 从它所在的堆删掉,再把 7 放进 large(它大于 small 堆顶)。此刻 small 有 1 个、large 有 2 个。
19. 21两堆个数失衡了,做一次再平衡:large 多了一个,把 large 堆顶 6 挪到 small,直到 small 比 large 多 0 或 1 个。现在窗口 [3, 6, 7] 的两堆排好了。
20. 22读 small 堆顶:6,就是窗口 [3, 6, 7] 的中位数 6。答案接成 [1, -1, -1, 3, 5, 6]。
21. 23六个窗口全部滑完,中位数依次是 [1, -1, -1, 3, 5, 6],和开头说的一致。回头看:每个窗口我们都没有重新排序,只靠「加一个、删一个、平衡一下、读堆顶」就拿到了中位数,每步只是常数次堆/有序集合操作(具体复杂度按实现区分,见下一页)。

⚠️ 容易写错的地方

完整代码（Python / C++ / Java）

from typing import List
import heapq
from collections import defaultdict
class Solution:
    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
        small, large = [], []
        delayed = defaultdict(int)
        small_size = large_size = 0
        def prune(heap):
            while heap:
                x = -heap[0] if heap is small else heap[0]
                if delayed[x]:
                    delayed[x] -= 1
                    heapq.heappop(heap)
                else:
                    break
        def balance():
            nonlocal small_size, large_size
            if small_size > large_size + 1:
                heapq.heappush(large, -heapq.heappop(small))
                small_size -= 1; large_size += 1
                prune(small)
            elif small_size < large_size:
                heapq.heappush(small, -heapq.heappop(large))
                small_size += 1; large_size -= 1
                prune(large)
        def add(x):
            nonlocal small_size, large_size
            if not small or x <= -small[0]:
                heapq.heappush(small, -x); small_size += 1
            else:
                heapq.heappush(large, x); large_size += 1
            balance()
        def remove(x):
            nonlocal small_size, large_size
            delayed[x] += 1
            if x <= -small[0]:
                small_size -= 1
                if x == -small[0]: prune(small)
            else:
                large_size -= 1
                if large and x == large[0]: prune(large)
            balance()
        def median():
            return float(-small[0]) if k % 2 else (-small[0] + large[0]) / 2.0
        ans = []
        for i, x in enumerate(nums):
            add(x)
            if i >= k:
                remove(nums[i - k])
            if i >= k - 1:
                ans.append(median())
        return ans

#include <set>
#include <vector>
using namespace std;
class Solution {
public:
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        multiset<long long> low, high;
        auto rebalance = [&]() {
            while (low.size() > high.size() + 1) { auto it = prev(low.end()); high.insert(*it); low.erase(it); }
            while (low.size() < high.size()) { auto it = high.begin(); low.insert(*it); high.erase(it); }
        };
        auto add = [&](long long x) {
            if (low.empty() || x <= *low.rbegin()) low.insert(x);
            else high.insert(x);
            rebalance();
        };
        auto remove = [&](long long x) {
            auto it = low.find(x);
            if (it != low.end()) low.erase(it);
            else { it = high.find(x); if (it != high.end()) high.erase(it); }
            rebalance();
        };
        vector<double> ans;
        for (int i = 0; i < (int)nums.size(); ++i) {
            add(nums[i]);
            if (i >= k) remove(nums[i-k]);
            if (i >= k - 1) ans.push_back(k % 2 ? (double)*low.rbegin() : ((double)*low.rbegin() + *high.begin()) / 2.0);
        }
        return ans;
    }
};

import java.util.*;
class Solution {
    TreeMap<Integer, Integer> low = new TreeMap<>(), high = new TreeMap<>();
    int lowSize = 0, highSize = 0, k;
    public double[] medianSlidingWindow(int[] nums, int k) {
        this.k = k; low.clear(); high.clear(); lowSize = highSize = 0;
        double[] ans = new double[nums.length - k + 1];
        int idx = 0;
        for (int i = 0; i < nums.length; i++) {
            add(nums[i]);
            if (i >= k) remove(nums[i - k]);
            if (i >= k - 1) ans[idx++] = median();
        }
        return ans;
    }
    private void add(int x) {
        if (lowSize == 0 || x <= low.lastKey()) { put(low, x, 1); lowSize++; }
        else { put(high, x, 1); highSize++; }
        balance();
    }
    private void remove(int x) {
        if (low.containsKey(x)) { put(low, x, -1); lowSize--; }
        else { put(high, x, -1); highSize--; }
        balance();
    }
    private void balance() {
        while (lowSize > highSize + 1) {
            int x = low.lastKey(); put(low, x, -1); lowSize--; put(high, x, 1); highSize++;
        }
        while (lowSize < highSize) {
            int x = high.firstKey(); put(high, x, -1); highSize--; put(low, x, 1); lowSize++;
        }
    }
    private double median() {
        return k % 2 == 1 ? low.lastKey() : ((double) low.lastKey() + high.firstKey()) / 2.0;
    }
    private void put(TreeMap<Integer,Integer> map, int key, int delta) {
        int val = map.getOrDefault(key, 0) + delta;
        if (val == 0) map.remove(key); else map.put(key, val);
    }
}

复杂度

面试官可能追问