Heap: Find median in a stream

Problem Statement:

You are given a data stream, you need to get the median of the elements after each integer is read.

Median can be found by:

If the data stream has odd numbers then middle will be considered as median.

If the data stream has even number, then the median will be arithmetic mean of the 2 middle values

Example:

Input: arr [1, 2, 3, 4, 5]

Output: [1, 1, 2, 2, 3]

Explanation:

After reading 1st element, 1 -> median = 1
After reading 2nd element, 1, 2 -> median = (1+2)/2 = 1
After reading 3rd element, 1, 2, 3 -> median = 2
After reading 4th element, 1, 2, 3, 4 -> median = (2+3)/2 = 2
After reading 5th element, 1, 2, 3, 4, 5 -> median = 3

Solution Explanation:

We will use heaps to solve the problem.

The intuition behind the heaps is that, we are interested only in the middle elements when all the elements are sorted.

We will break the array into 2 parts.

We need the largest element of the first array and smallest element of the second array.

Time Complexity: O(1)
Space Complexity: O(1)

Code Solution

#include <iostream>
#include <vector>
#include <algorithm>
#include <queue>

using namespace std;

class MedianFinder {
public:
    priority_queue<int, vector<int>, greater<int> > minHeap;
	priority_queue<int> maxHeap;
    MedianFinder() {}
    
    void addNum(int num) 
    {
        maxHeap.push(num);

        if(!minHeap.empty() && !maxHeap.empty() && maxHeap.top() > minHeap.top()) 
        {
            minHeap.push(maxHeap.top());
            maxHeap.pop();
        }

        if(maxHeap.size() > minHeap.size()+1) 
        {
            minHeap.push(maxHeap.top());
            maxHeap.pop();
        }

        if(minHeap.size() > maxHeap.size()+1) 
        {
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }
    }
    
    double findMedian() 
    {
        if(minHeap.size() > maxHeap.size()) return minHeap.top();
        if(maxHeap.size() > minHeap.size()) return maxHeap.top();
        else return ( minHeap.top() + maxHeap.top())/2.0;
    }
};

int main() 
{
    MedianFinder medianFinder;
    medianFinder.addNum(1);
    medianFinder.addNum(2);
    cout << medianFinder.findMedian() << endl; 
    medianFinder.addNum(3);
    cout << medianFinder.findMedian() << endl; 
    return 0;
}

Output

1.5
2