Two-Heap Median Maintenance

Description

Numbers arrive one at a time, and at any moment someone asks for the median. Re-sorting on every question is O(n log n) per query; a sorted list with insertion is O(n) per insert because elements shift. Two heaps give O(log n) inserts and O(1) medians.

This skill covers the two-heap split: a max-heap for the smaller half and a min-heap for the larger half, plus the two invariants that keep the median at their roots. The same two-heap shape also separates "locked" items from "available" ones, as in the IPO problem.

We'll go from simple to more technical:

1. The pattern: lower half in a max-heap, upper half in a min-heap, rebalanced after each insert.
2. A second flavor: locked and available items (IPO), with a greedy pick from the available heap.
3. Where running medians show up (live p50 latency, smoothing, schedulers).
4. What breaks it: the wrong half, drifting sizes, overflow in the average, deletions.

The approach

PriorityQueue<Integer> lower = new PriorityQueue<>(Comparator.reverseOrder());  // max-heap: smaller half
PriorityQueue<Integer> upper = new PriorityQueue<>();                           // min-heap: larger half

void addNum(int num) {
    lower.offer(num);
    upper.offer(lower.poll());              // the largest of the lower half moves up
    if (upper.size() > lower.size()) {
        lower.offer(upper.poll());          // keep lower equal or one larger
    }
}

double findMedian() {
    if (lower.size() > upper.size()) return lower.peek();
    return ((long) lower.peek() + upper.peek()) / 2.0;   // widen before adding: MAX + MAX overflows an int
}

import heapq


class MedianFinder:
    def __init__(self):
        self.lower = []                          # max-heap of the smaller half, stored NEGATED
        self.upper = []                          # min-heap of the larger half

    def add_num(self, num: int) -> None:
        heapq.heappush(self.lower, -num)
        heapq.heappush(self.upper, -heapq.heappop(self.lower))   # the largest of the lower half moves up
        if len(self.upper) > len(self.lower):
            heapq.heappush(self.lower, -heapq.heappop(self.upper))   # keep lower equal or one larger

    def find_median(self) -> float:
        if len(self.lower) > len(self.upper):
            return -self.lower[0]
        return (-self.lower[0] + self.upper[0]) / 2  # Python ints can't overflow

Two invariants make the roots the median: every value in lower is at most every value in upper, and lower has the same size as upper or one more. Routing each number through lower and passing its maximum up keeps the first invariant without comparing against both roots; the rebalance keeps the second. The average needs care: lower.peek() + upper.peek() adds two Integers as ints, and (MAX_VALUE + MAX_VALUE) / 2.0 comes out as -1.0. Python's ints can't overflow; its trap is the max-heap - heapq has only min-heap functions before 3.14, so lower stores -num and every read negates it back.

A second flavor: two heaps as "locked" and "available"

Two heaps can also separate items by whether they're usable yet. IPO: you may finish at most k projects; each needs capital[i] to start and adds profits[i] when done. Keep projects sorted by capital (the locked ones), move every project you can now afford into a max-heap of profits, and take the most profitable:

int findMaximizedCapital(int k, int w, int[] profits, int[] capital) {
    Integer[] byCapital = /* indices 0..n-1 sorted by capital */;
    PriorityQueue<Integer> affordable = new PriorityQueue<>(Comparator.reverseOrder());
    int next = 0;
    for (int done = 0; done < k; done++) {
        while (next < byCapital.length && capital[byCapital[next]] <= w) affordable.add(profits[byCapital[next++]]);
        if (affordable.isEmpty()) break;    // nothing affordable: capital can't grow
        w += affordable.poll();
    }
    return w;
}

import heapq


def find_maximized_capital(k: int, w: int, profits: list[int], capital: list[int]) -> int:
    by_capital = sorted(zip(capital, profits))
    affordable = []                              # max-heap of profits, stored negated
    nxt = 0
    for _ in range(k):
        while nxt < len(by_capital) and by_capital[nxt][0] <= w:
            heapq.heappush(affordable, -by_capital[nxt][1])
            nxt += 1
        if not affordable:
            break                                # nothing affordable: capital can't grow
        w -= heapq.heappop(affordable)
    return w

Capital never shrinks, so an affordable project stays affordable; taking the most profitable one now is never worse than saving it (an exchange argument). Taking the least profitable affordable project first ends at 7 instead of 106 with capital 1, k = 2 and projects (profit 5, needs 1), (profit 100, needs 6), (profit 1, needs 0).

Where you'll see it

Live dashboards showing the median (p50) latency of a stream, running medians in signal smoothing, and schedulers that hold "not yet eligible" and "ready" jobs in two priority queues (by release time and by priority) all use the two-heap split.

Which tool for which job?

Common mistakes

Check yourself

Try to answer before opening each one - recalling beats re-reading.

1. What two invariants make the heap roots give the median?

2. Why route every new number through the lower heap and move its maximum up?

3. Why does (lower.peek() + upper.peek()) / 2.0 fail for large values?

Key takeaways

• Lower half in a max-heap, upper half in a min-heap; lower's size equals upper's or is one more.
• Insert through lower, move its max up, rebalance - O(log n); the median is O(1).
• Widen before averaging the roots.
• Two heaps also split "locked" from "available" items (IPO).