中位数是有序序列最中间的那个数。如果序列的大小是偶数,则没有最中间的数;此时中位数是最中间的两个数的平均数。
例如:
- [ 2 , 3 , 4 ] [2,3,4] [2,3,4],中位数是 3 3 3
- [ 2 , 3 ] [2,3] [2,3],中位数是 ( 2 + 3 ) / 2 = 2.5 (2 + 3) / 2 = 2.5 (2+3)/2=2.5
给你一个数组
n
u
m
s
nums
nums,有一个大小为
k
k
k 的窗口从最左端滑动到最右端。窗口中有
k
k
k 个数,每次窗口向右移动
1
1
1 位。你的任务是找出每次窗口移动后得到的新窗口中元素的中位数,并输出由它们组成的数组。
示例:
给出 n u m s = [ 1 , 3 , − 1 , − 3 , 5 , 3 , 6 , 7 ] nums = [1,3,-1,-3,5,3,6,7] nums=[1,3,−1,−3,5,3,6,7],以及 k = 3 k = 3 k=3。
窗口位置 中位数
--------------- -----
[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
因此,返回该滑动窗口的中位数数组 [ 1 , − 1 , − 1 , 3 , 5 , 6 ] [1,-1,-1,3,5,6] [1,−1,−1,3,5,6]。
我们首先思考一下完成本题需要做哪些事情:
- 初始时,我们需要将数组 nums \textit{nums} nums 中的前 k k k 个元素放入一个滑动窗口,并且求出它们的中位数;
- 随后滑动窗口会向右进行移动。每一次移动后,会将一个新的元素放入滑动窗口,并且将一个旧的元素移出滑动窗口,最后再求出它们的中位数。
因此,我们需要设计一个「数据结构」,用来维护滑动窗口,并且需要提供如下的三个接口:
-
insert(num) \texttt{insert(num)} insert(num):将一个数 num \textit{num} num 加入数据结构;
-
erase(num) \texttt{erase(num)} erase(num):将一个数 num \textit{num} num 移出数据结构;
-
getMedian() \texttt{getMedian()} getMedian():返回当前数据结构中所有数的中位数。
方法一:双优先队列 + 延迟删除
我们可以使用两个优先队列(堆)维护所有的元素,第一个优先队列 small \textit{small} small 是一个大根堆,它负责维护所有元素中较小的那一半;第二个优先队列 large \textit{large} large 是一个小根堆,它负责维护所有元素中较大的那一半。具体地,如果当前需要维护的元素个数为 x x x,那么 small \textit{small} small 中维护了 ⌈ x 2 ⌉ \lceil \frac{x}{2} \rceil ⌈2x⌉个元素, large \textit{large} large中维护了 ⌊ x 2 ⌋ \lfloor \frac{x}{2} \rfloor ⌊2x⌋个元素,其中 ⌈ y ⌉ \lceil y \rceil ⌈y⌉和 ⌊ y ⌋ \lfloor y \rfloor ⌊y⌋ 分别表示将 y y y 向上取整和向下取整。也就是说:
small \textit{small} small 中的元素个数要么与 large \textit{large} large 中的元素个数相同,要么比 large \textit{large} large 中的元素个数恰好多 1 1 1 个。
这样设计的好处在于:当二者包含的元素个数相同时,它们各自的堆顶元素的平均值即为中位数;而当 small \textit{small} small包含的元素多了一个时, small \textit{small} small 的堆顶元素即为中位数。这样 getMedian() \texttt{getMedian()} getMedian()就设计完成了。
而对于 insert(num) \texttt{insert(num)} insert(num)而言,如果当前两个优先队列都为空,那么根据元素个数的要求,我们必须将这个元素加入 small \textit{small} small;如果 small \textit{small} small 非空(显然不会存在 small \textit{small} small 空而 large \textit{large} large 非空的情况),我们就可以将 num \textit{num} num 与 small \textit{small} small 的堆顶元素 top \textit{top} top 比较:
-
如果 num ≤ top \textit{num} \leq \textit{top} num≤top,我们就将其加入 small \textit{small} small 中;
-
如果 num > top \textit{num} > \textit{top} num>top,我们就将其加入 large \textit{large} large 中。
在成功地加入元素 num \textit{num} num之后,两个优先队列的元素个数可能会变得不符合要求。由于我们只加入了一个元素,那么不符合要求的情况只能是下面的二者之一:
-
small \textit{small} small 比 large \textit{large} large 的元素个数多了 2 2 2 个;
-
small \textit{small} small 比 large \textit{large} large的元素个数少了 1 1 1 个。
对于第一种情况,我们将 small \textit{small} small 的堆顶元素放入 large \textit{large} large;对于第二种情况,我们将 large \textit{large} large 的堆顶元素放入 small \textit{small} small,这样就可以解决问题了, insert(num) \texttt{insert(num)} insert(num)也就设计完成了。
然而对于 erase(num) \texttt{erase(num)} erase(num)而言,设计起来就不是那么容易了,因为我们知道,优先队列是不支持移出非堆顶元素这一操作的,因此我们可以考虑使用「延迟删除」的技巧,即:
当我们需要移出优先队列中的某个元素时,我们只将这个删除操作「记录」下来,而不去真的删除这个元素。当这个元素出现在 small \textit{small} small 或者 large \textit{large} large的堆顶时,我们再去将其移出对应的优先队列。
「延迟删除」使用到的辅助数据结构一般为哈希表 delayed \textit{delayed} delayed,其中的每个键值对 ( num \textit{num} num, freq \textit{freq} freq),表示元素 num \textit{num} num还需要被删除 freq \textit{freq} freq 次。「优先队列 + 延迟删除」有非常多种设计方式,体现在「延迟删除」的时机选择。在本题解中,我们使用一种比较容易编写代码的设计方式,即:
我们保证在任意操作 insert(num) \texttt{insert(num)} insert(num), erase(num) \texttt{erase(num)} erase(num), getMedian() \texttt{getMedian()} getMedian()完成之后(或者说任意操作开始之前), small \textit{small} small和 large \textit{large} large的堆顶元素都是不需要被「延迟删除」的。这样设计的好处在于:我们无需更改 getMedian() \texttt{getMedian()} getMedian()的设计,只需要略加修改 insert(num) \texttt{insert(num)} insert(num)即可。
我们首先设计一个辅助函数 prune(heap) \texttt{prune(heap)} prune(heap),它的作用很简单,就是对 heap \textit{heap} heap这个优先队列( small \textit{small} small 或者 large \textit{large} large 之一),不断地弹出其需要被删除的堆顶元素,并且减少 delayed \textit{delayed} delayed 中对应项的值。在 prune(heap) \texttt{prune(heap)} prune(heap)完成之后,我们就可以保证 heap \textit{heap} heap的堆顶元素是不需要被「延迟删除」的。
这样我们就可以在 prune(heap) \texttt{prune(heap)} prune(heap) 的基础上设计另一个辅助函数 makeBalance() \texttt{makeBalance()} makeBalance(),它的作用即为调整 small \textit{small} small 和 large \textit{large} large中的元素个数,使得二者的元素个数满足要求。由于有了 erase(num) \texttt{erase(num)} erase(num)以及「延迟删除」,我们在将一个优先队列的堆顶元素放入另一个优先队列时,第一个优先队列的堆顶元素可能是需要删除的。因此我们就可以用 makeBalance() \texttt{makeBalance()} makeBalance() 将 prune(heap) \texttt{prune(heap)} prune(heap)封装起来,它的逻辑如下:
-
如果 small \textit{small} small 和 large \textit{large} largelarge 中的元素个数满足要求,则不进行任何操作;
-
如果 small \textit{small} small 比 large \textit{large} large 的元素个数多了 2 2 2个,那么我们我们将 small \textit{small} small 的堆顶元素放入 large \textit{large} large。此时 small \textit{small} small 的对应元素可能是需要删除的,因此我们调用 prune(small) \texttt{prune(small)} prune(small);
-
如果 small \textit{small} small比 large \textit{large} large的元素个数少了 1 1 1个,那么我们将 large \textit{large} large 的堆顶元素放入 small \textit{small} small。此时 large \textit{large} large 的对应的元素可能是需要删除的,因此我们调用 prune(large) \texttt{prune(large)} prune(large)。
此时,我们只需要在原先 insert(num) \texttt{insert(num)} insert(num) 的设计的最后加上一步 makeBalance() \texttt{makeBalance()} makeBalance() 即可。然而对于 erase(num) \texttt{erase(num)} erase(num),我们还是需要进行一些思考的:
-
如果 num \textit{num} num与 small \textit{small} small 和 large \textit{large} large的堆顶元素都不相同,那么 num \textit{num} num是需要被「延迟删除」的,我们将其在哈希表中的值增加 1 1 1;
-
否则,例如 num \textit{num} num 与 small \textit{small} small 的堆顶元素相同,那么该元素是可以理解被删除的。虽然我们没有实现「立即删除」这个辅助函数,但只要我们将 num \textit{num} num 在哈希表中的值增加 1 1 1,并且调用「延迟删除」的辅助函数 prune(small) \texttt{prune(small)} prune(small),那么就相当于实现了「立即删除」的功能。
无论是「立即删除」还是「延迟删除」,其中一个优先队列中的元素个数发生了变化(减少了 1 1 1),因此我们还需要用 makeBalance() \texttt{makeBalance()} makeBalance() 调整元素的个数。
此时,所有的接口都已经设计完成了。由于 insert(num) \texttt{insert(num)} insert(num)和 erase(num) \texttt{erase(num)} erase(num)的最后一步都是 makeBalance() \texttt{makeBalance()} makeBalance(),而 makeBalance() \texttt{makeBalance()} makeBalance() 的最后一步是 prune(heap) \texttt{prune(heap)} prune(heap),因此我们就保证了任意操作完成之后, small \textit{small} small 和 large \textit{large} large 的堆顶元素都是不需要被「延迟删除」的。
复杂度分析
由于「延迟删除」的存在, small \textit{small} small 比 large \textit{large} large在最坏情况下可能包含所有的 n n n 个元素,即没有一个元素被真正删除了。因此优先队列的大小是 O ( n ) O(n) O(n) 而不是 O ( k ) O(k) O(k) 的,其中 nn 是数组 nums \textit{nums} nums 的长度。
时间复杂度: O ( n log n ) O(n\log n) O(nlogn)。 insert(num) \texttt{insert(num)} insert(num)和 erase(num) \texttt{erase(num)} erase(num)的单次时间复杂度为 O ( log n ) O(\log n) O(logn), getMedian() \texttt{getMedian()} getMedian() 的单次时间复杂度为 O ( 1 ) O(1) O(1)。因此总时间复杂度为 O ( n log n ) O(n\log n) O(nlogn)。
空间复杂度: O ( n ) O(n) O(n)。即为 small , large 和 delayed \textit{small},\textit{large} 和 \textit{delayed} small,large和delayed 需要使用的空间。
#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
priority_queue<int> small;
priority_queue<int, vector<int>, greater<int>> large;
int smallSize = 0, largeSize = 0;
unordered_map<int, int> ma;
template<typename T>
void prune(T &heap) {
while (!heap.empty()) {
int x = heap.top();
if (!ma.count(x)) break;
ma[x]--;
if (ma[x] == 0) {
ma.erase(x);
}
heap.pop();
}
}
void makeBalance() {
if (smallSize >= largeSize + 2) {
large.push(small.top());
small.pop();
smallSize--;
largeSize++;
prune(small);
} else if (smallSize < largeSize) {
small.push(large.top());
large.pop();
smallSize++;
largeSize--;
prune(large);
}
}
void insert(int x) {
if (small.empty() || x <= small.top()) {
small.push(x);
smallSize++;
} else {
large.push(x);
largeSize++;
}
makeBalance();
}
void Delete(int x) {
ma[x]++;
if (x <= small.top()) {
smallSize--;
if (x == small.top()) {
prune(small);
}
} else {
largeSize--;
if (x == large.top()) {
prune(large);
}
}
makeBalance();
}
double getMedian(int k) {
return k & 1 ? small.top() : ((double) small.top() + large.top()) / 2;
}
vector<double> medianSlidingWindow(vector<int> &nums, int k) {
int n = nums.size();
for (int i = 0; i < k; i++) {
insert(nums[i]);
}
vector<double> res = {getMedian(k)};
for (int i = k; i < n; i++) {
insert(nums[i]);
Delete(nums[i - k]);
res.push_back(getMedian(k));
}
return res;
}
};
结语
读者可以尝试回答如下的两个问题来检验自己是否掌握了该方法:
在 insert(num) \texttt{insert(num)} insert(num)的最后我们加上了一步 makeBalance() \texttt{makeBalance()} makeBalance(),其中包括可能进行的 prune(heap) \texttt{prune(heap)} prune(heap) 操作,这对于 insert(num) \texttt{insert(num)} insert(num)操作而言是否是必要的?
在 insert(num) \texttt{insert(num)} insert(num) 的过程中,如果我们将 insert(num) \texttt{insert(num)} insert(num)放入了 large \textit{large} large 中,并且 num \textit{num} num 恰好出现在 large \textit{large} large 的堆顶位置,且两个优先队列的元素个数满足要求,不需要进行调整。此时会不会出现 num \textit{num} num是一个需要被「延迟删除」的元素的情况,这样就不满足在 insert(num) \texttt{insert(num)} insert(num)操作完成之后 large \textit{large} large 的堆顶是不需要被「延迟删除」的要求了?
答案
实际上是不必要的。因为在 insert(num) \texttt{insert(num)} insert(num)操作之前,两个优先队列的堆顶元素都是不需要被删除的,而我们只可能从那个被加入了一个元素的优先队列的堆顶元素放入另一个优先队列中,因此两个优先队列的堆顶元素仍然都是不需要被删除的。这样写只是为了将 insert(num) \texttt{insert(num)} insert(num) 和 erase(num) \texttt{erase(num)} erase(num)操作统一起来,减少代码的冗余。
不可能会出现这种情况,假设出现了这种情况,那么 num \textit{num} num 显然不会等于 large \textit{large} large 原先的堆顶元素,因为 large \textit{large} large 原先的堆顶元素一定是不需要被删除的。那么 num \textit{num} num 满足:
small
\textit{small}
small ~的堆顶元素 <
num
\textit{num}
num <
large
\textit{large}
large ~的堆顶元素
由于
small
\textit{small}
small 是大根堆,
large
\textit{large}
large 是小根堆,因此根本就不存在与
num
\textit{num}
num 值相同的元素,也就不可能会被延迟删除了。
class Solution {
public:
vector<double> medianSlidingWindow(vector<int> &nums, int k) {
int n = nums.size();
multiset<int> s;
vector<double> res;
for (int i = 0; i < n; i++) {
if (s.size() >= k) s.erase(s.find(nums[i - k]));
s.insert(nums[i]);
if (i >= k - 1) {
auto mid = s.begin();
advance(mid, k / 2);
res.push_back((0.0+*mid + *prev(mid, (1 - k % 2))) / 2);
}
}
return res;
}
};