Friday, December 9, 2011

Numbers are randomly generated and passed to a method. Write a program to find and maintain the median value as new values are generated.

One solution is to use two priority heaps: a max heap for the values below the median, and a min heap for the values above the median. The median will be largest value of the max heap.
When a new value arrives it is placed in the below heap if the value is less than or equal to the median, otherwise it is placed into the above heap. The heap sizes can be equal or the below heap has one extra. This constraint can easily be restored by shifting an element from one heap to the other. The median is available in constant time, so updates are O(lg n).

1       private Comparator<Integer> maxHeapComparator, minHeapComparator;
2       private PriorityQueue<Integer> maxHeap, minHeap;
3       public void addNewNumber(int randomNumber) {
4              if (maxHeap.size() == minHeap.size()) {
5                      if ((minHeap.peek() != null) &&
6                               randomNumber > minHeap.peek()) {
7                               maxHeap.offer(minHeap.poll());
8                               minHeap.offer(randomNumber);
9                      } else {
10                              maxHeap.offer(randomNumber);
11                     }
12             }
13             else {
14                     if(randomNumber < maxHeap.peek()){
15                              minHeap.offer(maxHeap.poll());
16                              maxHeap.offer(randomNumber);
17                     }
18                     else {
19                              minHeap.offer(randomNumber);
20                     }
21             }
22      }
23      public static double getMedian() {
24             if (maxHeap.isEmpty()) return minHeap.peek();
25             else if (minHeap.isEmpty()) return maxHeap.peek();
26             if (maxHeap.size() == minHeap.size()) {
27                     return (minHeap.peek() + maxHeap.peek()) / 2;
28             } else if (maxHeap.size() > minHeap.size()) {
29                     return maxHeap.peek();
30             } else {
31                     return minHeap.peek();
32             }
33      }

No comments:

Post a Comment