Median in a stream of Integers

The problem, simply

Numbers come in one at a time, like a stream. After each new number, say the median of all numbers so far.

The median is the middle number when you sort them. If there is an even count, it is the average of the two middle numbers. For the stream 5, 15, 1, 3, 2, 8 the medians after each step are 5, 10, 5, 4, 3, 4.

A story to remember it

Picture it

Examples

Way 1: the simple way

Keep all numbers in a sorted list. For each new number, find its spot with binary search and insert it there. The median is at the middle index.

1. Binary search for the place of the new number.
2. Insert it there. Everything after it shifts one step right.
3. If the count is odd, return the middle item. If even, return the average of the two middle items.

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

class MedianFinder {
    private final List<Integer> sorted = new ArrayList<>();

    public void addNum(int num) {
        int pos = Collections.binarySearch(sorted, num);
        if (pos < 0) pos = -(pos + 1);        // where it should go
        sorted.add(pos, num);                 // shifts everything after it
    }

    public double findMedian() {
        int n = sorted.size();
        if (n % 2 == 1) return sorted.get(n / 2);
        return ((long) sorted.get(n / 2 - 1) + sorted.get(n / 2)) / 2.0;
    }
}

How long it takes: the search is O(log n), but the insert shifts up to n items, so each add is O(n) and n adds are O(n^2). The median is O(1). Extra room: the list, O(n).

Way 2: the clever way

Use the two teams. The small half lives in a max heap (biggest of the small numbers on top). The big half lives in a min heap (smallest of the big numbers on top).

1. Add: if the small half is empty or num <= small.peek(), add to the small half. Otherwise add to the big half.
2. Balance: if the small half has more than one extra, move its top to the big half. If the big half is bigger, move its top to the small half.
3. Median: if the small half is bigger, return its top. Otherwise return the average of both tops.

import java.util.Collections;
import java.util.PriorityQueue;

class MedianFinder {
    private final PriorityQueue<Integer> small = new PriorityQueue<>(Collections.reverseOrder()); // max heap
    private final PriorityQueue<Integer> big = new PriorityQueue<>();                             // min heap

    public void addNum(int num) {
        if (small.isEmpty() || num <= small.peek()) small.offer(num);
        else big.offer(num);
        // keep small equal to big, or one bigger
        if (small.size() > big.size() + 1) big.offer(small.poll());
        else if (big.size() > small.size()) small.offer(big.poll());
    }

    public double findMedian() {
        if (small.size() > big.size()) return small.peek();
        return ((long) small.peek() + big.peek()) / 2.0;   // long avoids overflow
    }
}

How long it takes: each add does at most two offers and one poll, O(log n). Each median only peeks, O(1). Extra room: both heaps together hold every number, O(n).

Watch out for

• Decide which half may hold the extra number and stick to it. Here the small half does, so an odd count reads small.peek().
• peek() on an empty heap returns null, and unboxing it throws. That is why the first number goes to the small half with the isEmpty() check.
• Rebalance after every add, not only sometimes. The size rule must hold before every median.
• Dividing by 2 instead of 2.0 drops the half: (3 + 4) / 2 is 3, not 3.5.

What interviewers ask next

• What if all numbers are between 0 and 100? Keep a count for each value, 101 buckets. Add is O(1) and the median walks the buckets, O(100).
• What if 99 percent are between 0 and 100? Use the 101 buckets plus two counters for numbers below and above. The median almost always lands in a bucket.
• Median of a sliding window of size k? The same two heaps, but numbers also leave. Use lazy deletion (a map of numbers to remove when they reach a top) or two TreeMap multisets. Each step is O(log k).
• The 90th percentile instead of the median? Keep the small half at 90 percent of the size instead of half.
• How would you handle billions of numbers that do not fit in memory? Approximate with a sketch such as t-digest, or bucket by value ranges and count.

Remember this

Grown-up words: the stream is an online input, the short team is a max heap of the lower half, the tall team is a min heap of the upper half, walking over is rebalancing, and the middle kid is the median.