Calculate median using priority queues

13,058

This should do it. Maintain two priority queues of the numbers greater and less than the median value. Shift values between the two queues such that they stay balanced, or close to balanced, and define the median based on the top values of the priority queues.

#include <queue>
#include <vector>
#include <functional>
#include <limits>
#include <iostream>

template<class T>
class RunningMedian {
 private:
  //Values greater than the median sorted so that smallest is on top
  std::priority_queue<T, std::vector<T>, std::greater<T> > upper;
  //Values smaller than the median sorted so that greatest is on top
  std::priority_queue<T, std::vector<T>, std::less<T> > lower;

 public:
  RunningMedian(){
    //Seed the queues
    upper.push(std::numeric_limits<T>::max());
    lower.push (std::numeric_limits<T>::min());
  }
  void push(T val){
    //Add number to the property queue

    //If number is greater than the least upper number, it is an upper number
    if(val>=upper.top())
      upper.push(val);
    else //Otherwise it is a lower number
      lower.push(val);

    //Rebalance
    //If the queues sizes differ by 2, they need to be rebalanced so that they
    //differ by 1.
    if(upper.size()-lower.size()==2){ //Upper queue is larger
      lower.push(upper.top());        //Move value from upper to lower
      upper.pop();                    //Remove same value from upper
    } else if(lower.size()-upper.size()==2){ //Lower queue is larger
      upper.push(lower.top());               //Move value from lower to upper
      lower.pop();                           //Remove same value
    }
  }

  T getMedian() const {
    if(upper.size()==lower.size())               //Upper and lower are same size
      return(upper.top()+lower.top())/((T)2.0);  //so median is average of least upper and greatest lower
    else if(upper.size()>lower.size())           //Upper size greater
      return upper.top();
    else                                         //Lower size greater
      return lower.top();
  }
};

int main(){
  RunningMedian<int> rm;
  rm.push(10);
  rm.push(3);
  rm.push(1);
  rm.push(20);
  rm.push(5);
  rm.push(8);
  rm.push(-1);
  std::cout<<rm.getMedian()<<std::endl; //Gives 5
}
Share:
13,058
Zechariah Schneider
Author by

Zechariah Schneider

Updated on June 04, 2022

Comments

  • Zechariah Schneider
    Zechariah Schneider almost 2 years

    I'm required to calculate the median. I've been told that the best way to do this is in this particular application is with a priority queue. I have no clue on how to proceed. I'd really appreciate any help.

  • Zechariah Schneider
    Zechariah Schneider about 9 years
    Makes sense. Thanks a bunch!
  • Richard
    Richard about 9 years
    My pleasure. If you feel this answered your question, feel free to accept the answer. Otherwise, let me know if any part of your question remains unanswered.