Heaps: Find the Running Median

  • + 0 comments

    Getting timeout and one wrong answer with this: Any suggestions?

    #include <bits/stdc++.h>
    
    using namespace std;
    
    
    void maxHeapify(vector<int> &arr, int iBegin)
    {
        int iLeft = 2 * iBegin + 1;
        int iRight = iLeft + 1;
        int iLargest = iBegin;
        
        if (iLeft < arr.size() && arr[iLeft] > arr[iLargest])
            iLargest = iLeft;
    
        if (iRight < arr.size() && arr[iRight] > arr[iLargest])
            iLargest = iRight;
    
        if (iLargest != iBegin)
        {
            swap(arr[iLargest], arr[iBegin]);
            maxHeapify(arr, iLargest);
        }
    }
    
    void minHeapify(vector<int>& arr, int iBegin)
    {
        int iLeft = 2 * iBegin + 1;
        int iRight = iLeft + 1;
        int iSmallest = iBegin;
        
        if (iLeft < arr.size() && arr[iLeft] < arr[iSmallest])
            iSmallest = iLeft;
    
        if (iRight < arr.size() && arr[iRight] < arr[iSmallest])
            iSmallest = iRight;
    
        if (iSmallest != iBegin)
        {
            swap(arr[iSmallest], arr[iBegin]);
            minHeapify(arr, iSmallest);
        }
    }
    
    
    void buildMinHeap(vector<int>& arr)
    {
        for (int i = arr.size() / 2 - 1; i >= 0; i--)
            minHeapify(arr, i);
    }
    
    void buildMaxHeap(vector<int>& arr)
    {
        for (int i = arr.size() / 2 - 1; i >= 0; i--)
            maxHeapify(arr, i);
    }
    
    void findRunningMedian(vector<int> arr)
    {
        //median even = (maxHeap[0] + minHeap[0] ) / 2
        //median odd = maxHeap[0]
        int i = 0;
        vector<int> minHeap;
        vector<int> maxHeap;
        minHeap.reserve(arr.size() / 2);
        maxHeap.reserve(arr.size() / 2);
    
        cout << fixed;
        cout << setprecision(1);
    
        while (i < arr.size())
        {
            double median = 0.0;
    
            // add odd element to maxHeap, to keep it always >= minHeap
            if (i % 2 == 0)
            {
                // add to minHeap if greater than its min element
                if (minHeap.size() && minHeap[0] < arr[i])
                {
                    //add to minHeap
                    minHeap.push_back(arr[i]);
    
                    // move the min element to maxHeap to maintain size
                    maxHeap.push_back(minHeap[0]);
                    minHeap.erase(minHeap.begin());
    
                    buildMaxHeap(maxHeap);
                    buildMinHeap(minHeap);
                }
                else
                {
                    maxHeap.push_back(arr[i]);
                    buildMaxHeap(maxHeap);
                }
    
                median = maxHeap[0];
            }
            else
            {
                // add to maxHeap if less than minHeap[0]
                if (minHeap.size() && minHeap[0] > arr[i])
                {
                    //add to minHeap
                    maxHeap.push_back(arr[i]);
    
                    // move the max element to minHeap, to maintain size
                    minHeap.push_back(maxHeap[0]);
                    maxHeap.erase(maxHeap.begin());
                    buildMaxHeap(maxHeap);
                    buildMinHeap(minHeap);
                }
                else
                {
                    minHeap.push_back(arr[i]);
    
                    // special case to handle 2 element array
                    if (i == 1 && minHeap[0] < maxHeap[0])
                    {
                        swap(minHeap[0], maxHeap[0]);
                    }
                    buildMinHeap(minHeap);
                }
    
                median = ((double)maxHeap[0] + (double)minHeap[0]) / 2;
            }
            cout << median << endl;
            i++;
        }
    }
    
    
    int main()
    {
        int n;
        cin >> n;
        cin.ignore(numeric_limits<streamsize>::max(), '\n');
    
        vector<int> a(n);
    
        for (int i = 0; i < n; i++) {
            int a_item;
            cin >> a_item;
            cin.ignore(numeric_limits<streamsize>::max(), '\n');
    
            a[i] = a_item;
        }
    
        findRunningMedian(a);
        return 0;
    }