C++ 95 ms single multiset O(n * log n)


  • 3
    V

    The reason to use two heaps is to have O(1) lookup for the median. It's is O(n) if we use multiset, but what if we reuse the median pointer when we can?
    The solution below still has the O(n^2) worst case run-time, and the average case run-time is O(n * log n). We can achieve O(n * log n) in the worst case if we make sure that the multiset comparator is stable.

    vector<double> medianSlidingWindow(vector<int>& nums, int k)
    {
        int size = nums.size(), median_pos = k - k / 2 - 1;
        vector<double> res(size - k + 1);
        multiset<int> s(nums.begin(), nums.begin() + k);
        auto it = next(s.begin(), median_pos);
    
        for (auto i = k; i <= size; ++i)
        {
            res[i - k] = ((double)*it + (k % 2 != 0 ? *it : *next(it))) / 2;
            if (i < size)
            {
                // magic numbers (instead of enum) for brevity. INT_MAX means to retrace the iterator from the beginning.
                int repos_it = INT_MAX; 
                if (k > 2)
                {
                    // if inserted or removed item equals to the current median, we need to retrace.
                    // we do not know which exact element will be removed/inserted, and we cannot compare multiset iterators.
                    // otherwise, we can keep or increment/decrement the current median iterator.
                    if ((nums[i - k] < *it && nums[i] < *it) || (nums[i - k] > *it && nums[i] > *it)) repos_it = 0;
                    else if (nums[i - k] < *it && nums[i] > *it) repos_it = 1; // advance forward.
                    else if (nums[i - k] > *it && nums[i] < *it) repos_it = -1; // advance backward.
                }
                s.insert(nums[i]);
                s.erase(s.find(nums[i - k]));
    
                if (repos_it == INT_MAX) it = next(s.begin(), median_pos);
                else if (repos_it == 1) ++it;
                else if (repos_it == -1) --it;
            }
        }
        return res;
    }
    

  • 1

    Thanks for next, didn't know that yet. Here's mine now, I think it's O(n log k). If the new/obsolete number equals the middle number, I insert/erase it right there. This way I have better control and never have to retrace the middle from scratch.

    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        multiset<int> window(nums.begin(), nums.begin() + k);
        auto mid = next(window.begin(), k / 2);
        vector<double> medians;
        for (int i=k; ; i++) {
    
            // Push the current median.
            medians.push_back((double(*mid) + *next(mid, k%2 - 1)) / 2);
    
            // If all done, return.
            if (i == nums.size())
                return medians;
                
            // Insert nums[i].
            if (nums[i] == *mid) {
                mid = window.insert(mid, nums[i]);
            } else {
                window.insert(nums[i]);
                if (nums[i] < *mid)
                    mid--;
            }
    
            // Erase nums[i-k].
            if (nums[i-k] == *mid) {
                mid = window.erase(mid);
            } else {
                window.erase(window.find(nums[i-k]));
                if (nums[i-k] < *mid)
                    mid++;
            }
        }
    }

  • 0
    V

    @StefanPochmann Yes, it looks like O (n log k) even in the worst case. Great idea to use the returned iterator from insert/erase. I did not know erase can return an iterator; it looks like it was added in C++ 11.

    I personally think the solution with a single multiset is cleaner/easier to understand than the "two heaps" one. Re-balancing and selecting which heap to insert/remove from could be tricky to get right.


  • 0

    @votrubac Another way to insert, without using the returned iterator:

            // Insert nums[i].
            if (nums[i] == *mid)
                window.insert(mid, nums[i]);
            else
                window.insert(nums[i]);
            if (nums[i] <= *mid)
                mid--;
    

    I'm sure erase could be done similarly...

    Btw, you keep saying "worth case", but surely you mean "worst case" :-)


  • 0

    Missed one "worth" in your edits :-P

    You could also do

            s.insert(s.lower_bound(nums[i]), nums[i]);
            s.erase(s.lower_bound(nums[i - k]));
    

    and then you would know which element will be removed/inserted (always the first with that value). Not sure whether it would help.


  • 1
    V

    @StefanPochmann We can simplify the insert even further, since in C++ 11
    "The relative ordering of equivalent elements is preserved, and newly inserted elements follow their equivalents already in the container."

    And thanks for spotting my "worth-s", not sure why my finders keep typing that...

        // Insert nums[i].
        window.insert(nums[i]);
        if (nums[i] < *mid) mid--;   
    

  • 0

    @votrubac Very nice! And I realized I can make erasing just as nice. Posted it as its own topic now.


Log in to reply
 

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