C++ 144ms solution, max and min heap using priority_queue


  • 1

    class MedianFinder {

    priority_queue<int> small;   // max heap
    priority_queue<int,vector<int>,greater<int>> big;  // min heap
    int szS,szB;
    
    public:
    MedianFinder(){szS=0;szB=0;}
    // Adds a number into the data structure.
    void keepbalance(){     // keep the different between szS and szB with 1.
        if (szS>szB+1){
            big.push(small.top());
            small.pop();
            szB++;szS--;
        }else if (szB>szS+1){
            small.push(big.top());
            big.pop();
            szB--;szS++;   
        }
    }
    void addNum(int num) {
        if (szS==0) {
            big.push(num); szB++;
        }else if (szB==0){
            small.push(num); szS++;
        }else{
            if (num<small.top()){
                small.push(num); szS++;
            }else{
                big.push(num); szB++;
            }
        }
        keepbalance();
    }
    
    // Returns the median of current data stream
    double findMedian() {
        if (szS==szB){
            return (small.top()+big.top())/2.0;
        }else if(szS>szB) return small.top();
        else return big.top();
    }
    };

  • 0
    W

    thanks for sharing. You don't have to store the size, just use big.size() is enough


Log in to reply
 

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