# 480 - Sliding Window Median

解法一 - Two Heaps

這題很直覺就會想要結合 sliding window 跟 Two Heaps,不過如果用 C++ 的 priority_queue,要刪除元素會滿麻煩的,所以可以用 set。

https://stackoverflow.com/a/55734733/1128197

概念上雖然簡單,不過實作上還是有一些難點:

  1. 可能有重複元素,所以要用 multiset,另外刪除時要刪除 iterator,不要刪除 value,不然可能會誤刪好幾個 element

  2. 縮減完 window 後才平衡兩個 heap 的 element 數,不然會不平衡

  3. 算 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