# 480 - Sliding Window Median
解法一 - Two Heaps
這題很直覺就會想要結合 sliding window 跟 Two Heaps,不過如果用 C++ 的 priority_queue,要刪除元素會滿麻煩的,所以可以用 set。
https://stackoverflow.com/a/55734733/1128197
概念上雖然簡單,不過實作上還是有一些難點:
可能有重複元素,所以要用 multiset,另外刪除時要刪除 iterator,不要刪除 value,不然可能會誤刪好幾個 element
縮減完 window 後才平衡兩個 heap 的 element 數,不然會不平衡
算 median 時注意 int overflow
程式碼實作如下:
class Solution {
public:
vector<double> medianSlidingWindow(vector<int>& nums, int k) {
vector<double> medians;
int windowStart = 0, n = nums.size();
for(int windowEnd = 0; windowEnd < n; windowEnd++) {
// Expand windowEnd and add to heaps
if(maxHeap.empty() or nums[windowEnd] <= *maxHeap.rbegin()) {
maxHeap.insert(nums[windowEnd]);
}
else {
minHeap.insert(nums[windowEnd]);
}
// If window size too big, move windowStart to reduce window
int windowSize = maxHeap.size() + minHeap.size();
if(windowSize > k) {
int toDelete = nums[windowStart++];
if(maxHeap.count(toDelete)) {
maxHeap.erase(maxHeap.find(toDelete));
}
else {
minHeap.erase(minHeap.find(toDelete));
}
}
// Balance maxHeap and minHeap
if(maxHeap.size() > minHeap.size() + 1) {
minHeap.insert(*maxHeap.rbegin());
auto it = maxHeap.end();
--it;
maxHeap.erase(it);
}
else if(minHeap.size() > maxHeap.size()) {
maxHeap.insert(*minHeap.begin());
minHeap.erase(minHeap.begin());
}
// Get median and put into medians
windowSize = maxHeap.size() + minHeap.size();
if(windowSize == k) {
if(k % 2 == 0) {
// prevent int overflow
long first = *minHeap.begin();
long second = *maxHeap.rbegin();
long long sum = first + second;
medians.push_back( sum / 2.0 );
}
else {
medians.push_back( *maxHeap.rbegin() );
}
}
/*
cout << "maxHeap contains: ";
for (set<int>::iterator it=maxHeap.begin(); it!=maxHeap.end(); ++it)
cout << ' ' << *it;
cout << "\nminHeap contains: ";
for (set<int>::iterator it=minHeap.begin(); it!=minHeap.end(); ++it)
cout << ' ' << *it;
cout << "\n================\n";
*/
}
return medians;
}
private:
multiset<int> maxHeap; // use rbegin() to get max value
multiset<int> minHeap; // use begin() to get min value
};
Last updated