Two short solutions


  • 0
    1. O(nlogk) The intuition is to reuse Find Median from Data Stream. Use two bst to store the smaller half and the larger half of the current window. After each add or erase, the two bst are balanced so that their sizes differ by at most 1. This makes it easier to partition the number and compute the median.
    class Solution {
    public:
        vector<double> medianSlidingWindow(vector<int>& nums, int k) {
            int n = nums.size();
            vector<double> med(n-k+1);
            for(int i=0;i<n;i++) {
                if(i>k-1) {
                    auto it = sm.find(nums[i-k]);
                    if(it!=sm.end()) sm.erase(it);
                    else lg.erase(lg.find(nums[i-k]));
                    balance();
                }
                if(sm.empty()||nums[i]<=*sm.rbegin()) sm.insert(nums[i]);
                else lg.insert(nums[i]);
                balance();
                if(i>=k-1) med[i-k+1] = sm.size()>lg.size() ? *sm.rbegin() : 
                    lg.size()>sm.size() ? *lg.begin() : ((long)*sm.rbegin()+*lg.begin())/2.0;
            }
            return med;
        }
    private:
        void balance() {
            if(sm.size()>lg.size()+1) {
                lg.insert(*sm.rbegin());
                sm.erase(--sm.end());
            }
            if(lg.size()>sm.size()+1) {
                sm.insert(*lg.begin());
                lg.erase(lg.begin());
            }
        }
        multiset<int> sm,lg;
    };
    
    1. O(nlogk) , we can just use one bst and keep track of the middle element. The great idea is from @StefanPochmann. The idea is to keep the middle pointer points to index k/2 after insertion and erasure.
        vector<double> medianSlidingWindow(vector<int>& nums, int k) {
            int n = nums.size();
            vector<double> med(n-k+1);
            multiset<double> window(nums.begin(),nums.begin()+k);
            auto m = next(window.begin(),k/2);
            for(int i=k;;i++) {
                med[i-k] = ((double)*m+*next(m,k%2-1))/2;
                if(i==n) return med;
                window.insert(nums[i]);
                if(nums[i]<*m) m--;
                if(nums[i-k]<=*m) m++;
                window.erase(window.lower_bound(nums[i-k]));
            }
        }
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.