295. Find Median from Data Stream

Description

Solution

Because we are handling a stream problem, if we need resort the array each time it is a costy solution. So if we could have a self-balance binary search tree, we could use the top of the tree as the answer, but just writing a AVL or RBT is not very troublesome. A better solution is to use two priority queue, one is min heap and other is max heap.

Code

class MedianFinder {
public:
    /** initialize your data structure here. */
    priority_queue<int,vector<int>, greater<>> heapMin;
    priority_queue<int,vector<int>, less<>> heapMax;
    MedianFinder() {
    }
    
    void addNum(int num) {
        if(heapMin.empty() || heapMin.top() <= num){
            heapMin.push(num);
            if(heapMin.size() > heapMax.size() + 1){
                heapMax.push(heapMin.top());
                heapMin.pop();
            }
        }else{
            heapMax.push(num);
            if(heapMax.size() >= heapMin.size() + 1){
                heapMin.push(heapMax.top());
                heapMax.pop();
            }
        }
    }
    
    double findMedian() {
        if(heapMax.size() == heapMin.size())
            return 1.0 * (heapMax.top() + heapMin.top()) /2;
        return heapMin.top();
    }
};