find_median

find_median.cpp

#include <queue>
#include <vector>
using namespace std;

class MedianFinder {
private:
  struct comp_min{bool operator()(int &a, int &b){return a > b;}};
  struct comp_max{bool operator()(int &a, int &b){return a < b;}};
  priority_queue<int, vector<int>, comp_min> rightQ; // [n/2, n]
  priority_queue<int, vector<int>, comp_max> leftQ; // [0, n/2-1]
  int size;
public:
    MedianFinder() {
      size = 0;
    }

    void addNum(int num) {
      if(size == 0){ rightQ.push(num); }
      else if(size == 1 || leftQ.size() < rightQ.size()){
        int leftTop = 0;
        int rightBot = rightQ.top();
        if(!leftQ.empty()) leftTop = leftQ.top();
        if(num > rightBot){
          rightQ.pop();
          rightQ.push(num);
          leftQ.push(rightBot);
        } else {
          leftQ.push(num);
        }
      }
      else{
        int leftTop = leftQ.top();
        int rightBot = rightQ.top();
        if(num < leftTop){
          leftQ.pop();
          leftQ.push(num);
          rightQ.push(leftTop);
        } else{
          rightQ.push(num);
        }
      }
      size++;
    }

    double findMedian() {
      int rightBot = rightQ.top();
      if(size == 1) return rightBot;
      int leftTop = leftQ.top();
      if(size%2 != 0) return rightBot;
      else return ((double)leftTop+(double)rightBot)/2;
    }
};