# Efficient data structure for finding median

Just come up with a relatively simple implementation for a data structure that can report median value in constant time, and insertion/remove median in `O(log(n))` time. It uses a max heap and min heap to store elements that are smaller and larger then the median respectively, and maintains the property that the differences between the two heaps are always less than or equal to 1.

Feel free to use/enhance/comment it.

```import java.util.Collections;
import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.Queue;

/**
* Data structure for storing collection of elements which can report
* new element and {@linkplain FastMedian#removeMedian() remove median} takes O(log(n)) time.
* If there are even number of elements, it always
* uses the smaller of the two median numbers as the median value.
*/
public class FastMedian<T> {
private final Comparator<? super T> comparator;
private final Queue<T> maxHeap;
private final Queue<T> minHeap;
private T median;

/**
* Creates a {@code FastMedian} which order elements according to their
*/
public FastMedian() {
this(null);
}

/**
* Creates a {@code FastMedian} which order elements according to the given
*
* @param comparator the comparator for ordering elements. If {@code null},
*        the {@linkplain Comparable natural ordering} will be used.
*/
public FastMedian(Comparator<? super T> comparator) {
this.comparator = comparator;
if (comparator == null) {
this.maxHeap = new PriorityQueue<T>(11, Collections.reverseOrder());
this.minHeap = new PriorityQueue<T>();
} else {
this.maxHeap = new PriorityQueue<T>(11, Collections.reverseOrder(comparator));
this.minHeap = new PriorityQueue<T>(11, comparator);
}
}

/**
* Inserts the given element.
*
* @param e the element to add.
*/
if (median == null) {
median = e;
return;
}
int cmp = (comparator == null) ? (((Comparable<? super T>) e).compareTo(median)) : comparator.compare(e, median);
if (cmp < 0) {
} else {
}
// Rebalance
int sizeDiff = maxHeap.size() - minHeap.size();
if (sizeDiff > 0) {
median = maxHeap.remove();
} else if (sizeDiff < -1) {
median = minHeap.remove();
}
}

/**
* Returns the current median value.
*
* @return current median value.
*/
public T getMedian() {
return median;
}

/**
* Removes the current median and returns it.
*
* @return value of the median removed.
*/
public T removeMedian() {
T result = median;
if (maxHeap.isEmpty() && minHeap.isEmpty()) {
median = null;
return result;
}
if (maxHeap.size() >= minHeap.size()) {
median = maxHeap.remove();
} else {
median = minHeap.remove();
}
return result;
}
}
```

A very simple example on using it:

```
import java.util.TreeSet;

public class MedianTest {

public void test() {
TreeSet<Integer> numbers = new TreeSet<Integer>();
FastMedian<Integer> fastMedian = new FastMedian<Integer>();

for (int i = 1; i < 10; i += 2) {
System.out.println(numbers + " median: " + fastMedian.getMedian());
}

for (int i = 10; i >= 2; i -= 2) {
System.out.println(numbers + " median: " + fastMedian.getMedian());
}

System.out.println();

Integer median = fastMedian.removeMedian();
while (median != null) {
numbers.remove(median);
System.out.println(numbers + " removed: " + median + " new median: " + fastMedian.getMedian());
median = fastMedian.removeMedian();
}
}
}
```

## 2 thoughts on “Efficient data structure for finding median”

1. chloepo

Hi there! could you tell me why is there a 11 inside this line code?
if (comparator == null) {
this.maxHeap = new PriorityQueue(11, Collections.reverseOrder());