Tuesday, October 20, 2015

[LeetCode] Find Median from Data Stream

Question can be found https://leetcode.com/problems/find-median-from-data-stream/

class MedianFinder {
    priority_queue<int> up;
    priority_queue<int, vector<int>, greater<int> > down;
public:
    // Adds a number into the data structure.
    void addNum(int num) {
        if(up.empty() || num <= up.top()) {
            up.push(num);
        }
        else if(down.empty() || num >= down.top()) {
            down.push(num);
        }
        else {
            up.push(num);
        }

        if(up.size() >= down.size() + 2) {
            int t = up.top();
            up.pop();
            down.push(t);
        }
        else if(down.size() >= up.size() + 2){
            int t = down.top();
            down.pop();
            up.push(t);
        }
    }
    // Returns the median of current data stream
    double findMedian() {
        if(up.size() > down.size()) return up.top();
        else if(down.size() > up.size()) return down.top();
        return (up.top() + down.top())/2.0;
    }
};