leetcode / Heap / priority queue / Two heaps advanced

Two heaps

The median is the middle of a sorted list, and keeping a list sorted while numbers arrive is expensive. But you never need the whole order, only the two items nearest the middle. Split the numbers into a lower half and an upper half, keep the largest of the lower half and the smallest of the upper half on top of two heaps, and the middle is always within reach in O(1), with O(log n) per arrival.

On this page
  1. Problem
  2. Intuition
  3. Mechanics
  4. The invariant
  5. Why not one heap?
  6. The unlock variant
  7. Worked example
  8. Implementation
  9. Variations and edge cases
  10. Complexity and tradeoffs
  11. Problems
  12. Hard
  13. Connections
  14. Read first
  15. Read next
  16. Check yourself
  17. References

Problem

Values arrive one at a time and after each arrival you must report the median of everything seen so far. Sorting on every query costs O(n log n) each time, and inserting into a sorted array costs O(n) each time. The same two-sided structure also solves a different kind of question: "unlock everything I can afford, then take the best one", where one heap orders what you can reach and another orders what you have not yet reached.

SignalPatternWhy
"Median of a stream" or "median after every insertion"Max-heap for the lower half, min-heap for the upper halfThe median is one or two heap roots and rebalancing is O(log n).
"Sliding-window median"Two heaps with lazy deletion, or a sorted containerThe same split, plus a way to forget items that left the window.
"Pick the best option among those currently affordable or available"Sorted list of unlock thresholds plus a max-heapThe pending items sit in sorted order and the unlocked ones sit in a heap.
"Keep the k-th element as items arrive and leave"Two heaps with a size ruleA lower half of fixed size k puts the k-th element at its top.
"Median of a fixed, static array"Sort or quickselectNo stream, so there is nothing to maintain; one O(n log n) or O(n) pass is enough.
"Values come from a small range"Counting array and a running totalIf values are bounded integers, counts find the median without heaps.

Intuition

Line everyone up by height and cut the line at the middle. You do not need to keep the people inside each half in order. You only need the tallest person in the short half and the shortest person in the tall half, because those two stand next to the cut. If the halves have equal size, the median is the average of those two; if one half has one more, the median is the extra half's edge person.

That pair is exactly what a heap exposes cheaply. The short half is a max-heap (its largest member is on top), and the tall half is a min-heap (its smallest member is on top). When a new number arrives, you send it through the cut: it may belong in either half, and the cut may no longer be in the middle, so a person is passed across to fix the sizes. Each fix moves at most one item, which is why it stays O(log n).

The two halves around the median After the values 5, 15, 1, 3, 8, 7, 9, 10 the lower half holds 7, 5, 3, 1, the largest of them, 7, at the cut. The upper half holds 8, 9, 10, 15, the smallest of them, 8, at the cut. A dashed line marks the cut, and the median is 7.5. 7 5 3 1 8 9 10 15 lower half: max-heap, top is its largest upper half: min-heap, top is its smallest

The heaps do not keep their members in this sorted order (a heap only orders the top), but the picture shows what they represent. Green and amber are the two roots next to the cut. The diagram scrolls sideways on a narrow screen.

Mechanics

The invariant

Two facts hold between operations. Order: every value in the lower heap is at most every value in the upper heap, which is the same as saying the lower root is at most the upper root. Balance: the lower heap has either the same number of values as the upper heap or exactly one more. Together they give the median directly: with equal sizes it is the mean of the two roots, otherwise it is the lower root.

To add a value x, first push x into the lower heap, then move the lower root to the upper heap. That two-step push makes the order fact hold whatever x was: the value that crosses is the largest of the lower half plus x, which is no larger than anything already above the cut. Then, if the upper heap now has more values than the lower heap, move its root back down. Each add is at most three heap operations, each O(log n), and the median read is O(1).

Why not one heap?

A single heap exposes only one end. A min-heap tells you the smallest, but the median is in the middle of the order, and popping your way there destroys the heap. Two heaps meeting at the median expose both neighbours of the cut at once. This "two halves meeting in the middle" shape also appears in the sliding-window median and in problems that ask for the k-th element of a changing set.

The median after each arrival One column per arrival: 5, 15, 1, 3, 8, 7, 9, 10. Under each is the median of everything seen so far: 5, 10, 5, 4, 5, 6, 7, 7.5. A dot in each column sits at the height of its median. 5 5 10 15 5 1 4 3 5 8 6 7 7 9 7.5 10 arrival (bottom) and the median so far (dot)

The median wobbles as extreme values arrive on either side, but it moves by at most a step through the order each time, which is why only the two roots ever need to be looked at.

The unlock variant

The same two-sided structure solves "maximize capital with k projects". One side is a sorted list of projects not yet reachable, ordered by the capital they need. The other side is a max-heap of projects you can afford, ordered by profit. Each round, move every newly affordable project from the sorted list into the heap, then take the best. Because capital only grows, a project once affordable stays affordable, so nothing is ever moved back.

Worked example

Feed the numbers 5, 15, 1, 3, 8, 7, 9, 10 into the structure. Each row shows the two halves after the arrival, largest first on the left and smallest first on the right, and the median.

ArrivalLower half (max-heap)Upper half (min-heap)Median
55empty5
1551510
15, 1155
33, 15, 154
85, 3, 18, 155
75, 3, 17, 8, 156
97, 5, 3, 18, 9, 157
107, 5, 3, 18, 9, 10, 157.5

After the 8 arrivals the median is 7.5, the mean of 7 and 8. After the third arrival the lower half had one more value than the upper half, so its root, 5, was the median on its own.

Unlock the best project. Start with 0 capital and pick at most 3 of five projects needing capital 0, 1, 1, 3, 6 with profits 1, 2, 3, 5, 9.

RoundCapital beforeProfits newly unlockedBest profit takenCapital after
10111
212, 334
34559

The final capital is 9. Taking the biggest profit among the affordable ones is always safe because profits only add to capital, and more capital never makes a project less affordable.

Implementation

import heapq


class MedianFinder:
    def __init__(self):
        self.lower = []      # max-heap, stored as negated values
        self.upper = []      # min-heap

    def add(self, x):
        heapq.heappush(self.lower, -x)
        heapq.heappush(self.upper, -heapq.heappop(self.lower))
        if len(self.upper) > len(self.lower):
            heapq.heappush(self.lower, -heapq.heappop(self.upper))

    def median(self):
        if len(self.lower) > len(self.upper):
            return float(-self.lower[0])
        return (-self.lower[0] + self.upper[0]) / 2


def max_capital(rounds, wealth, profits, capital):
    projects = sorted(zip(capital, profits))
    available = []                           # max-heap of profits, negated
    i = 0
    for _ in range(rounds):
        while i < len(projects) and projects[i][0] <= wealth:
            heapq.heappush(available, -projects[i][1])
            i += 1
        if not available:
            break
        wealth -= heapq.heappop(available)
    return wealth

Python has only a min-heap, so the lower half stores negated values; negate again when reading. Return a float, or divide by two carefully, if the problem wants decimal medians.

Variations and edge cases

  • Sliding window. Use the same halves, but also record each value that leaves the window and delete it lazily: when it reaches a root, pop it. Track the logical sizes of the halves separately from the stored sizes, since stale entries inflate the heaps.
  • Other percentiles. Keep the lower heap at a chosen fraction of the total instead of half, and its root is that percentile.
  • Bounded values. If every value is an integer in a small range, a counting array with a running prefix finds the median in O(range), which beats the heaps when the range is small.
  • Order of operations. Pushing to the lower heap first, then moving across, is not the only ordering, but it handles every case with no comparison with the roots. The "compare with the root and choose a side" version is faster by a constant and easier to get wrong.
  • Edge cases: asking for the median before any value arrives (guard it), a single value, duplicates (they split across the halves freely), negative numbers (the negation trick is still fine), and an even count (average two roots, in floating point).

Complexity and tradeoffs

  • Time: O(log n) per add (at most three heap operations) and O(1) per median query.
  • Space: O(n) for the stored values, split between the two heaps.
  • Baseline. Insert with bisect.insort into a sorted list and read the middle: O(log n) to find the spot but O(n) to shift, with a tiny constant. For up to tens of thousands of values it is often as fast in practice and much shorter to write. It is a fair baseline, not a trick.
  • When it fails. With deletions in the middle (sliding window) plain heaps need lazy deletion, which is fiddly. A balanced tree with subtree sizes, or two sorted containers, handles removals directly. For one-off queries on a static array, quickselect finds the median in expected O(n).

Problems

The two home problems are the two uses of the pattern: a running median, where the heaps hold the lower and upper halves, and the capital-maximizing project picker, where one side is a sorted queue of what is out of reach.

Hard

  • 295. Find Median from Data Stream Blind 75 Grind 169 NeetCode 150 NeetCode 250
    Keep a max-heap of the lower half and a min-heap of the upper half balanced in size, so the median sits on top of them.
  • 502. IPO NeetCode 250
    Sort projects by required capital, move each newly affordable one into a profit max-heap, and always start the most profitable.

Connections

Read first

Read next

No page builds on this one yet.

This page builds on the heap operations from the heaps page. The size-k heap of the top-k page keeps one end of the order; here two heaps keep both neighbours of a cut. Quickselect, next, finds the same middle element without keeping any structure, trading the stream for an in-place pass over a fixed array. The greedy pages later revisit the "take the best of what is currently available" idea behind the project picker.

Check yourself

Why is the lower half a max-heap and the upper half a min-heap?

The median depends on the values next to the cut: the largest of the lower half and the smallest of the upper half. A max-heap puts the largest of its members on top and a min-heap the smallest, so both neighbours are available in O(1). In the example they were 7 and 8.

Why push to the lower heap and then move its root to the upper heap?

Whatever the new value is, the value that crosses to the upper heap is the largest of the lower half plus the newcomer, which is no larger than anything already in the upper half. That keeps the order fact true without comparing against either root.

What does the median equal when the heaps have different sizes?

The lower heap is allowed to have one extra value, and that extra value is the middle one, so the median is the lower root. After the third arrival of the example the halves held 2 and 1 values and the median was 5.

In the project picker, why can we never regret unlocking a project early?

Capital never decreases, because every profit is positive or zero, so a project that is affordable now is affordable in every later round. Moving it into the profit heap early loses nothing, and the heap lets us take the best one when we need it.

References

  • Cormen, Leiserson, Rivest and Stein, Introduction to Algorithms, chapter 6 (heaps and priority queues) and chapter 9 (medians and order statistics).
  • Sedgewick and Wayne, Algorithms, section 2.4, which includes the median-maintenance exercise with two priority queues.
  • Skiena, The Algorithm Design Manual, the sections on priority queues and on selection.
  • Python documentation, heapq, and the bisect module for the sorted-list alternative.
  • The NeetCode roadmap and the Blind 75, Grind 169, NeetCode 150 and NeetCode 250 lists, used as the index of topics and problems.