题目
https://leetcode-cn.com/problems/sliding-window-median/
双优先队列加延迟删除
由于要得到一个窗口内的中位数,可以使用两个优先队列small(大顶堆)、large(小顶堆),一个维护窗口内元数排序后较小的一半,一个维护较大的一半,如果元素个数为奇数那个中位数为small的堆顶元素,如果元素个数为偶数那么中位数就是small与large的堆顶元素除以2,此题的难点在于怎样保持两个堆的元素数量各占一半。
一开始将元素加入small中,如果small不为空则后续元素与small堆顶元素比较,如果小于small的堆顶元素则加入small中,如果大于small的堆顶元素则加入到large中。为了维持small与large的元素数量关系,当small的元素数量比large多2个时,就将small的堆顶元素加入到large中,如果small的元素数量比large少一个就将large的堆顶元素加入到small中。
由于滑动窗口需要删除元素,而堆不能删除中间位置的元素,因此采用延迟删除的方法,将要删除的元素加入哈希表中,如果该元素到堆顶时再把该元素删除。因此使用辅助函数prune(heap)来不断删除堆顶元素,保证了堆顶元素一定不是要删除的元素,这样就方便计算中位数了。
class Solution {
public double[] medianSlidingWindow(int[] nums, int k) {
DualHeap dh = new DualHeap(k);
for (int i = 0; i < k; i++) {
dh.insert(nums[i]);
}
double[] ans = new double[nums.length - k + 1];
ans[0] = dh.getMediean();
for (int i = k; i < nums.length; i++) {
dh.insert(nums[i]);
dh.erase(nums[i - k]);
ans[i - k + 1] = dh.getMediean();
}
return ans;
}
class DualHeap {
// 大根堆,维护较小的一般元素
private PriorityQueue<Integer> small;
// 小根堆,维护较大的一般元素
private PriorityQueue<Integer> large;
// 哈希表,记录被延迟删除的元素 key为要删除的元素,value为需要删除的次数
private Map<Integer, Integer> delayed;
// 滑动窗口大小
private int k;
// 由于延迟删除的操作,所以需要两个变量维护 small 和 large 的实际大小
private int smallSize, largeSize;
public DualHeap(int k) {
this.k = k;
small = new PriorityQueue<Integer>((o1, o2) -> o2.compareTo(o1));
large = new PriorityQueue<Integer>();
delayed = new HashMap<>();
smallSize = 0;
largeSize = 0;
}
public double getMediean() {
return (k & 1) == 1 ? small.peek() : ((double)small.peek() + large.peek()) / 2;
}
public void insert(int num) {
if (small.isEmpty() || num < small.peek()) {
small.offer(num);
smallSize++;
} else {
large.offer(num);
largeSize++;
}
makeBalance();
}
public void erase(int num) {
delayed.put(num, delayed.getOrDefault(num, 0) + 1);
if (num <= small.peek()) {
--smallSize;
if (num == small.peek()) {
prune(small);
}
} else {
--largeSize;
if (num == large.peek()) {
prune(large);
}
}
makeBalance();
}
private void prune(PriorityQueue<Integer> heap) {
while (!heap.isEmpty()) {
int num = heap.peek();
if (delayed.containsKey(num)) {
delayed.put(num, delayed.get(num) - 1);
if (delayed.get(num) == 0) {
delayed.remove(num);
}
heap.poll();
} else {
break;
}
}
}
private void makeBalance() {
if (smallSize > largeSize + 1) {
large.offer(small.poll());
largeSize++;
smallSize--;
prune(small);
} else if (smallSize < largeSize) {
small.offer(large.poll());
largeSize--;
smallSize++;
prune(large);
}
}
}
}