Find Median from Data Stream

median.cpp

#include <queue>
#include <vector>
#include <functional>

using namespace std;

class MedianFinder {
public:
	priority_queue<int, vector<int>> maxHeap; //store the numbers which are less than the median
	priority_queue<int, vector<int>, greater<int>> minHeap; // numbers greater than the median

	
	// Adds a number into the data structure.
	void addNum(int num) {
		cout << maxHeap.size() << " " << minHeap.size() << endl;
		if (minHeap.size() == 0 || num > minHeap.top()){
			minHeap.push(num);
			if (minHeap.size() - maxHeap.size() > 1){  //do balance
				maxHeap.push(minHeap.top());
				minHeap.pop();
			}
		}
		else{
			maxHeap.push(num);
			if (maxHeap.size() - minHeap.size() > 1){
				minHeap.push(maxHeap.top());
				maxHeap.pop();
			}
		}		
	}

	double findMedian() {
		if (maxHeap.size() == minHeap.size()){
			return ((double)maxHeap.top() + minHeap.top()) / 2;
		}
		else if (maxHeap.size() > minHeap.size()){
			return maxHeap.top();
		}
		else{
			return minHeap.top();
		}
	}


};