CODEWITHSUNDEEP
javapriority-queuemedian
flash
flash

Reputation: 1519

optimize adding number to queue for median from stream of numbers

I was recently asked this question in an interview to find the median from data stream of numbers and I was able to come up with Priority Queue solution as shown below:

public class MedianFinder {
  private final PriorityQueue<Long> min = new PriorityQueue<>();
  private final PriorityQueue<Long> max = new PriorityQueue<>(Collections.reverseOrder());

  public void addNum(long num) {
    max.offer(num);
    min.offer(max.poll());
    if (max.size() < min.size()) {
      max.offer(min.poll());
    }
  }

  public double findMedian() {
    if (max.size() == min.size())
      return (max.peek() + min.peek()) / 2.0;
    else
      return max.peek();
  }
}

Now interviewer wanted me to optimize addNum method because it has lot of O(log n) operations (around 5) and he wanted to see if we can optimize it any further so that we have fewer O(log n) operations? Is there anything we can do here to optimize addNum method?

Upvotes: 2

Views: 92

Answers (1)

Ricky Mo
Ricky Mo

Reputation: 7648

This can reduce the average number of offer call from 2.5 to 1.5 and poll call from 1.5 to 0.5. Overall reduce the average number of O(log n) operations from 4 to 2.

public void addNum(long num) {
    if(!max.isEmpty() )
    {
        if(max.size() == min.size())
        {
            if(num > max.peek())
            {
                min.offer(num);
                max.offer(min.poll());
            }
            else
            {
                max.offer(num);
            }
        }
        else
        {
            if(num > max.peek())
            {
                min.offer(num);
            }
            else
            {
                max.offer(num);
                min.offer(max.poll());
            }
        }
    }
    else
    {
        max.offer(num);
    }
}

A more compact version (same logic)

public void addNum(long num) {
    if(!max.isEmpty())
    {
        (num > max.peek() ? min : max).offer(num);
        if(min.size() > max.size())
        {
            max.offer(min.poll());
        }
        else if(max.size() - min.size() > 1)
        {
            min.offer(max.poll());
        }
    }
    else
    {
        max.offer(num);
    }
}

Upvotes: 2

Related Questions