480. 滑动窗口中位数
中位数是有序序列最中间的那个数。如果序列的大小是偶数,则没有最中间的数;此时中位数是最中间的两个数的平均数。
例如:
[2,3,4]
,中位数是3
[2,3]
,中位数是(2 + 3) / 2 = 2.5
给你一个数组 nums
,有一个大小为 k
的窗口从最左端滑动到最右端。窗口中有 k
个数,每次窗口向右移动 1
位。你的任务是找出每次窗口移动后得到的新窗口中元素的中位数,并输出由它们组成的数组。
示例:
给出 nums = [1,3,-1,-3,5,3,6,7]
,以及 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]。
提示:
- 你可以假设 k 始终有效,即:k 始终小于输入的非空数组的元素个数。
- 与真实值误差在 10 -5 以内的答案将被视作正确答案。
方法一:优先队列
解题思路
此题的暴力解法是很容易想出来的:维护一个长度为 k 的数组,左出右进,每次排序后取中位数即可,但这样的时间复杂度是 O(klogk)
。在此基础上,思考优化解法。
我们知道 优先队列 中的元素是有序的,但是优先队列并不支持按索引取值。为此,可以维护两个优先队列 smallHeap
和 largeHeap
,各自存一半的数据,如果 k 为奇数 那么smallHeap
多放一个。注意:
smallHeap
是维护较小的那一半数据,堆顶是最大值,所以smallHeap
是一个「大顶堆」largeHeap
是维护较大的那一半数据,堆顶是最小值,所以largeHeap
是一个「小顶堆」- k 是奇数,中位数就是
smallHeap.peek()
;反之,中位数就是(smallHeap.peek() + largeHeap.peek()) / 2
存储的问题搞定了,还有删除。优先队列只能删除堆顶元素。为此可以额外使用一个 Map 维护需要删除的元素,key 为待删除的数字,value 为此数字需要删除的次数。当堆顶元素在 Map 中时,说明此时需要删除堆顶的元素。
此外,我们还需要保持两个队列的平衡,即 smallHeap.size()
要么等于 largeHeap.size()
,要么等于 largeHeap.size() + 1
。因为两个队列都可能存在 “待删除元素”,所以不能直接使用 size,需要手动维护两个队列中 “真实元素” 的数量。
参考代码
- 代码中封装了
DualHeap
类,其中#insert(int x)
、#remove(int x)
、#getMedian()
分别代表 插入元素、删除元素、获取中位数。 - 每次插入元素和删除元素后,都需要调用
#balance()
方法平衡两个队列以及#checkTop()
方法检查堆顶元素是否可以删除。
public double[] medianSlidingWindow(int[] nums, int k) {
int n = nums.length;
double[] ans = new double[n - k + 1];
DualHeap dualHeap = new DualHeap((k & 1) == 1);
for (int i = 0; i < n; i++) {
dualHeap.insert(nums[i]);
if (i < k - 1) {
continue;
}
ans[i - k + 1] = dualHeap.getMedian();
dualHeap.remove(nums[i - k + 1]);
}
return ans;
}
@SuppressWarnings("all")
class DualHeap {
// k 是否为奇数
private boolean odd;
private PriorityQueue<Integer> smallHeap;
private PriorityQueue<Integer> largeHeap;
private Map<Integer, Integer> lazyMap;
// 两个队列中真实的元素数量(去掉待删除的)
int smallSize, largeSize;
public DualHeap(boolean odd) {
this.odd = odd;
this.smallHeap = new PriorityQueue<>(Comparator.reverseOrder());
this.largeHeap = new PriorityQueue<>(Comparator.naturalOrder());
this.lazyMap = new HashMap<>();
smallSize = 0; largeSize = 0;
}
public void insert(int val) {
if (smallHeap.isEmpty() || val <= smallHeap.peek()) {
smallSize++;
smallHeap.offer(val);
} else {
largeSize++;
largeHeap.offer(val);
}
balance();
checkTop();
}
public void remove(int val) {
if (val <= smallHeap.peek()) {
smallSize--;
} else {
largeSize--;
}
lazyMap.put(val, lazyMap.getOrDefault(val, 0) + 1);
balance();
checkTop();
}
public double getMedian() {
return odd ? smallHeap.peek() : ((double) smallHeap.peek() + largeHeap.peek()) / 2;
}
// 执行插入和删除操作后,需要维持两个队列的平衡
private void balance() {
if (smallSize > largeSize + 1) {
smallSize--;
largeSize++;
largeHeap.offer(smallHeap.poll());
} else if (smallSize < largeSize) {
smallSize++;
largeSize--;
smallHeap.offer(largeHeap.poll());
}
}
// 如果堆顶元素是待删除元素则删除
private void checkTop() {
while (!smallHeap.isEmpty() && lazyMap.containsKey(smallHeap.peek())) {
int x = smallHeap.poll();
int count = lazyMap.get(x);
if (count == 1) {
lazyMap.remove(x);
} else {
lazyMap.put(x, count - 1);
}
}
while (!largeHeap.isEmpty() && lazyMap.containsKey(largeHeap.peek())) {
int x = largeHeap.poll();
int count = lazyMap.get(x);
if (count == 1) {
lazyMap.remove(x);
} else {
lazyMap.put(x, count - 1);
}
}
}
}
执行结果