C++ Solution O(n*k)


  • 5
    K

    The idea is to maintain a BST of the window and just search for the k/2 largest element and k/2 smallest element then the average of these two is the median of the window.

    Now if the STL's multiset BST maintained how many element were in each subtree finding each median would take O(log k) time but since it doesn't it takes O(k) time to find each median.

    public:
        vector<double> medianSlidingWindow(vector<int>& nums, int k) {
            multiset<int> mp;
            vector<double> med;
            
            for(int i=0; i<k-1; ++i) mp.insert(nums[i]);
            
            for(int i=k-1; i< nums.size(); ++i){
                mp.insert(nums[i]); // Add the next number
                
                auto itb = mp.begin(); advance(itb, (k-1)/2); //Find the lower median
                auto ite = mp.end(); advance(ite, -(k+1)/2); //Find the upper median
            
                double avg = ((long)(*itb) + (*ite)) / 2.0;
                med.push_back(avg);
            
                mp.erase(mp.find(nums[i-k+1])); //Remove the oldest element
            }
            
            return med;
        }
    };

  • 1
    S

    There are two advance in the for loop, you can optimize one of them to the following , and almost reduce the run time half.

        vector<double> medianSlidingWindow(vector<int>& nums, int k) {
            multiset<int> window;
            vector<double> ret;
            for (int i = 0; i < nums.size(); i++) {
                window.insert(nums[i]);
                if (window.size() < k)  continue;
                auto it1 = window.begin(), it2 = it1;
                advance(it1, (k-1)/2);
                it2 = it1;
                advance(it2, (k & 1) == 0);
                ret.push_back((long(*it1) + *it2) / 2.0);
                window.erase(window.find(nums[i-k+1]));
            }
            return ret;
        }
    

  • 0
    K

    @saxion Thanks for pointing out that optimization.


Log in to reply
 

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