Two Heaps

Two Heaps

Some problems need quick access to the middle of a changing collection: the running median of a stream of numbers, or the median of every window of k numbers. Keeping the data fully sorted costs O(n) per insertion. The two heaps pattern gets the middle in O(1), with O(log n) insertions, by keeping two halves:

• a max-heap holding the smaller half, whose top is the largest of the small numbers;
• a min-heap holding the larger half, whose top is the smallest of the large numbers.

The median sits right at the boundary between the two tops.

   small half (max-heap)         large half (min-heap)
   [ 1, 3, 4, 7 ]  top = 7   |   top = 9  [ 9, 12, 15 ]
                            median = 7 (odd total: the bigger half's top)

The two invariants

Everything depends on keeping these true after every operation:

1. Order: every number in small ≤ every number in large, i.e. small.peek() <= large.peek().
2. Balance: the sizes are equal, or small has exactly one more.

With both in place, the median is small.peek() when the total is odd, and the average of both tops when it's even.

Find Median from Data Stream

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

    public void addNum(int num) {
        // 1. Push through `small` so the largest small value moves to `large` → keeps the ORDER invariant
        small.offer(num);
        large.offer(small.poll());
        // 2. Restore BALANCE: `small` may be one bigger, never smaller
        if (large.size() > small.size()) small.offer(large.poll());
    }

    public double findMedian() {
        return small.size() > large.size()
                ? small.peek()
                : ((long) small.peek() + large.peek()) / 2.0;   // long avoids int overflow when summing
    }
}

Trace it: add 5, so small=[5], large=[]. Add 2, so small=[2], large=[5], and the median is 3.5. Add 9, so small=[5,2], large=[9], and the median is 5.

• addNum: O(log n), a constant number of heap operations. findMedian: O(1).
• Routing every number through small first is what makes the ordering automatic. You never have to compare the new number with both tops yourself.

Follow-up variants interviewers like: if all numbers are in the range 0–100, use a count array of 101 buckets and walk it to the middle: O(1) insert, O(100) median. If 99% of numbers are in 0–100, keep counters for the range plus two small overflow counters (below/above), since the median almost certainly lies inside the range.

Sliding Window Median

Now numbers also leave. PriorityQueue.remove(Object) is O(k), which gives O(n·k) overall. Two efficient approaches:

Approach 1: two heaps with lazy deletion. Don't remove a leaving number immediately. Record it in a "to delete" map, keep track of each heap's valid size, and discard invalid numbers only when they reach the top of a heap.

public double[] medianSlidingWindow(int[] nums, int k) {
    PriorityQueue<Integer> small = new PriorityQueue<>(Comparator.reverseOrder());
    PriorityQueue<Integer> large = new PriorityQueue<>();
    Map<Integer, Integer> delayed = new HashMap<>();     // value → how many copies are waiting to be deleted
    int[] size = new int[2];                             // valid sizes: [0] = small, [1] = large
    double[] result = new double[nums.length - k + 1];

    for (int i = 0; i < nums.length; i++) {
        // insert
        if (small.isEmpty() || nums[i] <= small.peek()) { small.offer(nums[i]); size[0]++; }
        else                                            { large.offer(nums[i]); size[1]++; }
        rebalance(small, large, size, delayed);

        if (i >= k) {                                    // remove the number leaving the window (lazily)
            int out = nums[i - k];
            delayed.merge(out, 1, Integer::sum);
            if (out <= small.peek()) { size[0]--; if (out == small.peek()) prune(small, delayed); }
            else                     { size[1]--; if (out == large.peek()) prune(large, delayed); }
            rebalance(small, large, size, delayed);
        }
        if (i >= k - 1) {
            result[i - k + 1] = (k % 2 == 1) ? small.peek() : ((double) small.peek() + large.peek()) / 2.0;
        }
    }
    return result;
}

private void rebalance(PriorityQueue<Integer> small, PriorityQueue<Integer> large, int[] size, Map<Integer, Integer> delayed) {
    if (size[0] > size[1] + 1)  { large.offer(small.poll()); size[0]--; size[1]++; prune(small, delayed); }
    else if (size[0] < size[1]) { small.offer(large.poll()); size[1]--; size[0]++; prune(large, delayed); }
}

private void prune(PriorityQueue<Integer> heap, Map<Integer, Integer> delayed) {
    while (!heap.isEmpty() && delayed.getOrDefault(heap.peek(), 0) > 0) {   // drop invalid tops
        delayed.merge(heap.poll(), -1, Integer::sum);
    }
}

Each number is inserted and popped at most once per heap, so the total is O(n log n).

Approach 2: two TreeMap multisets (value → count), which support O(log k) removal of any value directly. The code is simpler to reason about than lazy deletion, and it's a fine interview answer.

IPO: two heaps for a greedy choice

Here the two heaps answer two different questions: "which projects have become affordable?" (a min-heap by required capital) and "of the affordable ones, which is best?" (a max-heap by profit).

public int findMaximizedCapital(int k, int w, int[] profits, int[] capital) {
    PriorityQueue<int[]> locked = new PriorityQueue<>(Comparator.comparingInt(p -> p[0]));        // by capital
    PriorityQueue<int[]> affordable = new PriorityQueue<>((a, b) -> Integer.compare(b[1], a[1]));   // by profit desc
    for (int i = 0; i < profits.length; i++) locked.offer(new int[]{capital[i], profits[i]});

    for (int round = 0; round < k; round++) {
        while (!locked.isEmpty() && locked.peek()[0] <= w) affordable.offer(locked.poll());   // unlock
        if (affordable.isEmpty()) break;                     // nothing affordable: stop early
        w += affordable.poll()[1];                           // greedily take the most profitable
    }
    return w;
}

Greedy is correct because capital only grows. Taking the most profitable affordable project never hurts future options. It's O(n log n).

Recognizing the pattern

Reach for two heaps when a problem asks for:

• the median (or any percentile boundary) of a stream or window;
• repeatedly choosing the best item among those that have become eligible (IPO, scheduling);
• the smallest of the large and the largest of the small at the same time.

Follow-up questions this topic invites — and their answers

Q: Why not keep one sorted list and read the middle? A: Inserting into a sorted array or list is O(n) because of shifting. Two heaps give O(log n) inserts and O(1) median reads. A balanced BST or skip list with order statistics also works in O(log n), but two heaps are simpler to write under interview pressure.

Q: Why does addNum push every number through the max-heap first? A: It guarantees the order invariant without extra comparisons: whatever is largest in the small half moves to the large half. The rebalance step then only has to fix sizes. Without this, a large number could land in the small heap and break the ordering.

Q: How do you avoid overflow when averaging the two middle values? A: Add them as long (or double) before dividing: ((long) a + b) / 2.0. With int, two values near Integer.MAX_VALUE overflow before the division happens.

Q: What's the complexity of the sliding window median, and why not use PriorityQueue.remove? A: With lazy deletion or TreeMaps it's O(n log k) (or O(n log n) with lazy heaps). PriorityQueue.remove(Object) must search the heap linearly, O(k) per window, giving O(n·k), which is too slow for large k.

Q: How would you find the 90th percentile instead of the median? A: Use the same two heaps, but change the balance rule so the low heap holds about 90% of the elements: rebalance so small.size() stays at ceil(0.9 × total). The boundary top is then the 90th percentile.