Solution: Two Heaps

The sorted list keeps the whole stream in order, but median only ever reads the one or two middle values. A heap gives $O (1)$ access to one value, the one at its root. The middle values sit at the boundary between two halves: the largest of the smaller half and the smallest of the larger half. One root cannot hold both, so we keep two heaps, one for each half:

The smaller half goes in a max-heap, called low. Its root is the largest of the small values, the value just below the boundary.
The larger half goes in a min-heap, called high. Its root is the smallest of the large values, the value just above the boundary.

As in the kth largest solutions, low is an ArrayHeap built with a reversed comparator, and high uses the natural ordering.

The invariants

Two rules keep the split correct:

The boundary rule: every value in low is less than or equal to every value in high. This is what makes the two roots the true middle of the data, rather than two arbitrary values.
The size rule: the two sizes differ by at most one. When the total count is odd, we let low hold the one extra element.

If both rules hold, we can read the median straight from the roots. When the sizes are equal, the count is even and the two middle values are exactly the two roots, so the median is their average. When the sizes differ, the count is odd and the single middle value is the root of the larger heap, which is always low.

public class TwoHeapsMedianFinder implements MedianFinder {

  private PriorityQueue<Integer> low;   // the smaller half, largest on top
  private PriorityQueue<Integer> high;  // the larger half, smallest on top

  public TwoHeapsMedianFinder() {
    Comparator<Integer> largestFirst = Comparator.reverseOrder();
    low = new ArrayHeap<>(largestFirst);
    high = new ArrayHeap<>();
  }

  @Override
  public void add(int value) {
    if (low.size() == 0 || value <= low.peek()) {
      low.add(value);
    } else {
      high.add(value);
    }
    // Keep the sizes within one of each other, with any extra value in low.
    if (low.size() > high.size() + 1) {
      int moved = low.peek();  // the largest of the smaller half
      low.remove();
      high.add(moved);
    } else if (high.size() > low.size()) {
      int moved = high.peek();  // the smallest of the larger half
      high.remove();
      low.add(moved);
    }
  }

  @Override
  public double median() {
    if (low.size() == 0) {
      throw new NoSuchElementException();
    }
    if (low.size() == high.size()) {
      return ((double) low.peek() + high.peek()) / 2;  // even count
    }
    return low.peek();  // odd count: low holds the extra value
  }
}

Two things to notice in add. First, we choose the heap by comparing against low.peek(), the current boundary. Second, after the value goes in, low may hold two more values than high, or high may hold one more value than low. Moving a single root across restores the size rule without breaking the boundary rule. remove returns nothing, so we read the root with peek before removing it.

In median, an empty low means no values have been added, because the first value always goes into low, and any extra value stays there.

Tracing the stream

Let’s run the same stream, 5, 2, 8, 1, 9, 3. After each add we show both heaps with their roots marked, and read the median. Each heap is drawn in its array order, root first.

add(5). The low heap is empty, so 5 goes into low. Nothing to rebalance.

   low  (max-heap)                    high (min-heap)
   ┌───┐                              (empty)
   │ 5 │
   └───┘
     ▲
   peek → 5  (max)

   sizes 1 and 0  →  odd count, median = low root = 5

add(2). Since 2 <= low.peek() (which is 5), it goes into low. Now low has two elements and high has none, breaking the size rule:

   low  (max-heap)                    high (min-heap)
   ┌───┬───┐                          (empty)
   │ 5 │ 2 │
   └───┴───┘

So we rebalance: low.peek() is 5, we remove it from low and add it to high.

   low  (max-heap)                    high (min-heap)
   ┌───┐                              ┌───┐
   │ 2 │                              │ 5 │
   └───┘                              └───┘
     ▲                                  ▲
   peek → 2  (max)                    peek → 5  (min)

   sizes 1 and 1  →  even count, median = (2 + 5) / 2 = 3.5

add(8). Now 8 <= low.peek() (which is 2) is false, so 8 goes into high. That makes high bigger than low:

   low  (max-heap)                    high (min-heap)
   ┌───┐                              ┌───┬───┐
   │ 2 │                              │ 5 │ 8 │
   └───┘                              └───┴───┘

We rebalance the other way: high.peek() is 5, we remove it from high and add it to low.

   low  (max-heap)                    high (min-heap)
   ┌───┬───┐                          ┌───┐
   │ 5 │ 2 │                          │ 8 │
   └───┴───┘                          └───┘
     ▲                                  ▲
   peek → 5  (max)                    peek → 8  (min)

   sizes 2 and 1  →  odd count, median = low root = 5

add(1). Since 1 <= low.peek() (which is 5), 1 goes into low. Now low has three elements and high has one, breaking the size rule:

   low  (max-heap)                    high (min-heap)
   ┌───┬───┬───┐                      ┌───┐
   │ 5 │ 2 │ 1 │                      │ 8 │
   └───┴───┴───┘                      └───┘

We move low’s root, 5, across to high:

   low  (max-heap)                    high (min-heap)
   ┌───┬───┐                          ┌───┬───┐
   │ 2 │ 1 │                          │ 5 │ 8 │
   └───┴───┘                          └───┴───┘
     ▲                                  ▲
   peek → 2  (max)                    peek → 5  (min)

   sizes 2 and 2  →  even count, median = (2 + 5) / 2 = 3.5

add(9). Here 9 <= low.peek() (which is 2) is false, so 9 goes into high. That makes high bigger than low:

   low  (max-heap)                    high (min-heap)
   ┌───┬───┐                          ┌───┬───┬───┐
   │ 2 │ 1 │                          │ 5 │ 8 │ 9 │
   └───┴───┘                          └───┴───┴───┘

We move high’s root, 5, across to low:

   low  (max-heap)                    high (min-heap)
   ┌───┬───┬───┐                      ┌───┬───┐
   │ 5 │ 1 │ 2 │                      │ 8 │ 9 │
   └───┴───┴───┘                      └───┴───┘
     ▲                                  ▲
   peek → 5  (max)                    peek → 8  (min)

   sizes 3 and 2  →  odd count, median = low root = 5

add(3). Since 3 <= low.peek() (which is 5), 3 goes into low. Now low has four elements and high has two:

   low  (max-heap)                    high (min-heap)
   ┌───┬───┬───┬───┐                  ┌───┬───┐
   │ 5 │ 3 │ 2 │ 1 │                  │ 8 │ 9 │
   └───┴───┴───┴───┘                  └───┴───┘

We move low’s root, 5, across to high:

   low  (max-heap)                    high (min-heap)
   ┌───┬───┬───┐                      ┌───┬───┬───┐
   │ 3 │ 1 │ 2 │                      │ 5 │ 9 │ 8 │
   └───┴───┴───┘                      └───┴───┴───┘
     ▲                                  ▲
   peek → 3  (max)                    peek → 5  (min)

   sizes 3 and 3  →  even count, median = (3 + 5) / 2 = 4

In order, the medians are 5, 3.5, 5, 3.5, 5, 4, the same answers the sorted list produced.

What it costs

Reading the median is peek on one or both heaps, which is $O (1)$ . Each add does a constant number of heap operations: one add, a peek, and at most one root move from one heap to the other, which is a peek, a remove, and an add. On a binary heap, add is $O (lo g n)$ amortized, remove is $O (lo g n)$ , and peek is $O (1)$ , so one insert into the stream is $O (lo g n)$ amortized. Over $n$ values, the whole run is $O (n lo g n)$ , faster than the sorted list’s $O (n^{2})$ .