Calculate the running median

Question

I'm trying to solve a challenge where you need to calculate the median every time you add a number.

say you have a list of numbers: 1, 2, 3, 4, 5, 6, 7, 8, 9, 10

When you scan in the first number (1), you calculate the median of that. When you scan in the second number (2), you calculate the median of 1 and 2, and so on.

My code is working, but is having timeout issues on the larger test cases:

import java.io.*;
import java.util.*;
public class Solution {
 public static void main(String[] args) {
 Scanner s = new Scanner(System.in);
 int n = s.nextInt();
 PriorityQueue<Integer> pq = new PriorityQueue<>();
 while(n-- > 0) {
 pq.add(s.nextInt());
 System.out.println(calculateMedian(pq));
 }
 }
 public static float calculateMedian(PriorityQueue<Integer> pq) {
 PriorityQueue<Integer> temp = new PriorityQueue<>(pq);
 int index = 0;
 float number = 0.0f;
 int halfQueueSize = (temp.size() >> 1);
 if(temp.size() % 2 != 0) {
 while(index++ <= halfQueueSize) {
 number = temp.poll();
 }
 return number;
 } else {
 while(index++ < halfQueueSize) {
 number = temp.poll();
 }
 return ((number+temp.poll())/2.0f);
 }
 }
}

Any tips as to how this can be optimized would be much appreciated!

• Steve Jessop: A remark on the algorithm rather than the code itself - you should maintain two priority queues, a min heap containing values larger than the current median, and a max heap containing values smaller than the current median. Then when a new value comes in you need to do at most two pushes and one pop.
• maaartinus: Where can I find the challenge? +++ @SteveJessop That's pretty smart. IIRC the time for insertion is \$O(log n)\$ so it won't be better than my TreeMap w.r.t. the complexity, but practically it should be way faster.
• Nilzone-: @maaartinus hackerrank.com/challenges/find-median-1

Answer 1

The PriorityQueue was no good idea. What are you doing?

After the nth element, you copy all the previous n-1 elements into a temporary queue, which has complexity \$O(n)\$. Then you drain half the queue, which has complexity \$O(n \log n)\$. Summed up you get \$O(n^2 \log n)\$.

You're just misusing the queue as a sorting vehicle. That's not optimal as Arrays.sort has a faster algorithm (quicksort is faster, except when it goes wrong and this one tries pretty hard to avoid going wrong).

It would be probably faster to use an int[], track the size manually and re-sort the array after each element added. This has the same complexity, but could be faster as you at least avoid the copying.

Inserting the element in the proper position would lead to a total complexity of \$O(n^2)\$.

A much faster solution (I hope \$O(n \log n)\$) could be implemented using a TreeSet with higher and lower. There are also smart algorithms allowing to compute median without sorting the collection as you sort of need to sort only the middle part.

Review

public static float calculateMedian(PriorityQueue<Integer> pq) {

Don't use float unless you really have to conserve memory. They're pretty imprecise and they're no faster than double, at least on modern Desktop CPUs.

 int halfQueueSize = (temp.size() >> 1);
 if(temp.size() % 2 != 0) {

Do you optimize division by 2 as shift or don't you? Note that it hardly matters as this is surely not the most time-consuming operation here.

 if(temp.size() % 2 != 0) {
 while(index++ <= halfQueueSize) {
 number = temp.poll();
 }
 return number;
 } else {
 while(index++ < halfQueueSize) {
 number = temp.poll();
 }
 return ((number+temp.poll())/2.0f);
 }

You really shouldn't have two nearly identical branches. It makes it hard to spot the difference. With something like

 int halfQueueSize = (temp.size() >> 1) + (temp.size() & 1);

you can use < in both branches and pull the loop out.

• Nilzone-: thanks for such a detailed answer. Only thing I can add is: (temp.size() >> 1) + (temp.size() % 2 != 0); won't work since you're adding a int and a boolean. But I guess this would be the same thing? (temp.size() >> 1) + (temp.size() & 1);?
• maaartinus: @Nilzone- This was a copy & paste error, thanks, fixed. I'm not claiming that the expression is right, but something like this will surely do.