C++ O(nlogn) sol using two heap storing indices


  • 0
    S

    Using the technique similar to median of data stream, storing small elements indices and large elements indices.
    During iteration, when the heap top elements have indices out of sliding window, discard those elements. Then after each iteration, the top elements of two heaps are both in the sliding window and two heaps have the same number of elements that are in the window.

    class Solution {
    public:
        vector<double> medianSlidingWindow(vector<int>& nums, int k) {
            auto gt = [&nums](int& i1, int& i2) {
                return nums[i1] > nums[i2];
            };
            auto lt = [&nums](int& i1, int& i2) {
                return nums[i1] < nums[i2];
            };
            priority_queue<int, vector<int>, decltype(lt)> maxHeap(lt);
            priority_queue<int, vector<int>, decltype(gt)> minHeap(gt);
            
            vector<double> res;
            for (int i = 0; i < nums.size(); i++) {
                if (i < k) {
                    // form the initial heap
                    if (i % 2 == 0) {
                        // push to maxHeap 1 element;
                        if (minHeap.empty() || nums[i] <= nums[minHeap.top()]) maxHeap.push(i);
                        else {
                            maxHeap.push(minHeap.top());
                            minHeap.pop();
                            minHeap.push(i);
                        }
                    } else {
                        // push to minHeap 1 element;
                        if (maxHeap.empty() || nums[i] >= nums[maxHeap.top()]) minHeap.push(i);
                        else {
                            minHeap.push(maxHeap.top());
                            maxHeap.pop();
                            maxHeap.push(i);
                        }
                    }
                } else {
                    if (maxHeap.top() == i - k) {
                        // maxHeap top element becomes out of window, discard it and add new element in maxHeap
                        maxHeap.pop();
                        while (!maxHeap.empty() && maxHeap.top() < i - k) maxHeap.pop();
                        
                        // push to maxHeap 1 element;
                        if (minHeap.empty() || nums[i] <= nums[minHeap.top()]) maxHeap.push(i);
                        else {
                            maxHeap.push(minHeap.top());
                            minHeap.pop();
                            minHeap.push(i);
                        }
                        while (!minHeap.empty() && minHeap.top() < i - k) minHeap.pop();
                    } else if (minHeap.top() == i - k) {
                      // minHeap top element becomes out of window, discard it and add new element in maxHeap
                        minHeap.pop();
                        while (!minHeap.empty() && minHeap.top() < i - k) minHeap.pop();
                        // push to minHeap 1 element;
                        if (maxHeap.empty() || nums[i] >= nums[maxHeap.top()]) minHeap.push(i);
                        else {
                            minHeap.push(maxHeap.top());
                            maxHeap.pop();
                            maxHeap.push(i);
                        }
                        while (!maxHeap.empty() && maxHeap.top() < i - k) maxHeap.pop();
                    } else if (nums[i - k] <= nums[maxHeap.top()]) {
                        // the element should be discarded is not on top of either heap, then add new element to heaps
    
                        // push to maxHeap 1 element;
                        if (minHeap.empty() || nums[i] <= nums[minHeap.top()]) maxHeap.push(i);
                        else {
                            maxHeap.push(minHeap.top());
                            minHeap.pop();
                            minHeap.push(i);
                        }
                        while (!minHeap.empty() && minHeap.top() < i - k) minHeap.pop();
                    } else if (nums[i - k] >= nums[minHeap.top()]) {
                        // push to minHeap 1 element;
                        if (maxHeap.empty() || nums[i] >= nums[maxHeap.top()]) minHeap.push(i);
                        else {
                            minHeap.push(maxHeap.top());
                            maxHeap.pop();
                            maxHeap.push(i);
                        }
                        while (!maxHeap.empty() && maxHeap.top() < i - k) maxHeap.pop();
                    }
                }
                
                
                if (i >= k - 1) {
                    if (k % 2 == 0) res.push_back( ((double) nums[maxHeap.top()] + (double) nums[minHeap.top()]) / 2.0 );
                    else            res.push_back(  (double) nums[maxHeap.top()] );
                }
                
            }
            return res;
        }
    };
    

Log in to reply
 

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