Skip to content

Two Heaps and Running Median

One heap gives you the extreme of a collection in O(1)O(1). Sometimes you need the middle instead — and a heap cannot give you that, because the middle is exactly the element a heap makes no promises about.

The trick is to stop asking one heap for the middle and instead split the data in half:

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

Those two tops sit either side of the median. Keep the halves balanced in size and the median is always readable in O(1)O(1), with O(logn)O(\log n) inserts.

What you’ll learn

  • Python has no max-heap — the negation trick, and where it bites.
  • The push-then-shift insert that keeps both heaps correctly ordered without any comparison logic.
  • The balance invariant, and why “differ by at most one” is not quite enough.
  • The second use of two heaps: one heap sorted by availability, another by value — which is what LC 502 actually needs.
  • Three real LeetCode problems solved in the browser: 295, 502, 1962.

The cue

The template — running median

The elegant part of this insert is that it needs no comparison to decide which heap a new value belongs in. Always push into lowlow, then immediately move lowlow’s maximum across to highhigh. That guarantees every value in lowlow is ≤ every value in highhigh, whatever order things arrive in.

median_finder.py
import heapq
 
 
class MedianFinder:
    def __init__(self):
        self.low = []      # max-heap (negated): the smaller half
        self.high = []     # min-heap: the larger half
        # invariant: len(low) == len(high)  or  len(low) == len(high) + 1
 
    def add_num(self, num):
        # 1. always push into low, 2. shift low's max into high
        heapq.heappush(self.low, -num)
        heapq.heappush(self.high, -heapq.heappop(self.low))
        # 3. rebalance so low is never smaller than high
        if len(self.high) > len(self.low):
            heapq.heappush(self.low, -heapq.heappop(self.high))
 
    def find_median(self):
        if len(self.low) > len(self.high):
            return float(-self.low[0])                  # odd count
        return (-self.low[0] + self.high[0]) / 2.0      # even count
 
 
mf = MedianFinder()
for x in [1, 2, 3, 4]:
    mf.add_num(x)
print(mf.find_median())    # 2.5
median_finder.py
import heapq
 
 
class MedianFinder:
    def __init__(self):
        self.low = []      # max-heap (negated): the smaller half
        self.high = []     # min-heap: the larger half
        # invariant: len(low) == len(high)  or  len(low) == len(high) + 1
 
    def add_num(self, num):
        # 1. always push into low, 2. shift low's max into high
        heapq.heappush(self.low, -num)
        heapq.heappush(self.high, -heapq.heappop(self.low))
        # 3. rebalance so low is never smaller than high
        if len(self.high) > len(self.low):
            heapq.heappush(self.low, -heapq.heappop(self.high))
 
    def find_median(self):
        if len(self.low) > len(self.high):
            return float(-self.low[0])                  # odd count
        return (-self.low[0] + self.high[0]) / 2.0      # even count
 
 
mf = MedianFinder()
for x in [1, 2, 3, 4]:
    mf.add_num(x)
print(mf.find_median())    # 2.5

Complexity

OperationTwo heapsSorted list (bisect.insortbisect.insort)Re-sort each query
InsertO(logn)O(\log n)O(n)O(n) (shifting)O(1)O(1)
Find medianO(1)O(1)O(1)O(1)O(nlogn)O(n \log n)
nn inserts + queriesO(nlogn)O(n \log n)O(n2)O(n^2)O(n2logn)O(n^2 \log n)

bisect.insortbisect.insort keeps a sorted list and reads the median in O(1)O(1), but each insert shifts O(n)O(n) elements. It is genuinely competitive for small nn because the shifting is a fast memmove — worth mentioning, but the two-heap answer is what the O(logn)O(\log n) requirement is asking for.

The variant map

VariantThe two heaps holdCanonical problem
Running medianSmaller half (max-heap) / larger half (min-heap)295 Find Median from Data Stream
Median over a windowSame, plus lazy deletion of expired values480 Sliding Window Median
Unlock then chooseSorted by cost/availability / by value502 IPO · 2542
Modify and re-insertOne heap suffices, but items return changed1962 Remove Stones
Balance two groupsThe two groups themselvesScheduling / load-balancing variants

Practice — real LeetCode problems

LC 295 — Find Median from Data Stream · Hard

Problem. Implement MedianFinderMedianFinder with addNum(num)addNum(num) and findMedian()findMedian(). The median of an even-length collection is the mean of the two middle values.

Constraints. -10^5 <= num <= 10^5-10^5 <= num <= 10^5, up to 5 * 10^45 * 10^4 calls, findMedianfindMedian is only called after at least one addNumaddNum.

Example. Add 11, add 22, findMedian()findMedian() gives 1.51.5; add 33, findMedian()findMedian() gives 2.02.0.

Editorial — approach, complexity, follow-ups

Split the data into a smaller half and a larger half. The median is then determined entirely by the two heap tops — the largest small value and the smallest large value.

Time O(logn)O(\log n) per addNumaddNum, O(1)O(1) per findMedianfindMedian. Space O(n)O(n).

Two invariants do all the work:

  1. Ordering: every value in lowlow is ≤ every value in highhigh. The push-then-shift insert maintains this without comparing anything.
  2. Balance: len(low)len(low) equals len(high)len(high) or exceeds it by exactly one. Allowing only lowlow to be the larger one makes the odd case unambiguous.

Note float(-self.low[0])float(-self.low[0]) in the odd branch: LeetCode expects a float, and returning a bare intint can fail strict comparisons. The even branch divides with // so it is already a float.

Follow-ups you should expect:

  • “What if 99% of numbers are in a small range?” LeetCode asks this directly. If values are bounded integers, counting sort over the range gives O(1)O(1) insert and O(range)O(\text{range}) median — or a Fenwick tree over the value range for O(logrange)O(\log \text{range}) both ways.
  • “Sliding window median (LC 480)?” Add lazy deletion, or use SortedListSortedList.
  • “Track an arbitrary percentile?” Same structure, but the balance ratio becomes p : 1-pp : 1-p instead of 1 : 11 : 1.
  • “Remove a number too?” Heaps cannot delete arbitrary elements — lazy deletion with a CounterCounter, or switch to a balanced BST / SortedListSortedList.

LC 502 — IPO · Hard

Problem. You can complete at most kk projects. Project ii requires capital[i]capital[i] capital to start and yields profits[i]profits[i], which is added to your capital. Starting with ww capital, maximise your final capital.

Constraints. 1 <= k <= 10^51 <= k <= 10^5, 0 <= w <= 10^90 <= w <= 10^9, 1 <= len(profits) <= 10^51 <= len(profits) <= 10^5, 0 <= profits[i] <= 10^40 <= profits[i] <= 10^4, 0 <= capital[i] <= 10^90 <= capital[i] <= 10^9.

Examples. k = 2, w = 0, profits = [1,2,3], capital = [0,1,1]k = 2, w = 0, profits = [1,2,3], capital = [0,1,1] gives 44 · k = 3, w = 0, profits = [1,2,3], capital = [0,1,2]k = 3, w = 0, profits = [1,2,3], capital = [0,1,2] gives 66

Editorial — approach, complexity, follow-ups

Two different orderings are in play, and conflating them is the trap:

  • Availability is by capitalcapital required — you unlock cheapest-first.
  • Desirability is by profitprofit — you always want the best available one.

Sorting handles the first (a pointer sweeps forward and never goes back, since capital only grows). A max-heap handles the second.

Time O(nlogn)O(n \log n) for the sort plus O((n+k)logn)O((n + k) \log n) for the heap operations. Space O(n)O(n).

Why greedy is correct: capital never decreases, so a project affordable now stays affordable. Taking the largest available profit therefore cannot reduce your future options — it can only enlarge them. There is no reason to save a project for later.

Two details worth stating:

  • breakbreak when nothing is available. Not merely an optimisation: if ww is below the cheapest requirement, no further profit can arrive, so continuing would loop kk times doing nothing. (1, 0, [1,2,3], [1,1,2])(1, 0, [1,2,3], [1,1,2]) returns 00 for exactly this reason — nothing is affordable at w = 0w = 0.
  • kk may exceed the project count. The breakbreak handles that too, which the fourth case (k = 10k = 10, three projects) checks.

The “unlock by one key, choose by another” shape is the transferable idea. It reappears in task scheduling, meeting rooms with priorities, and LC 2542.

Follow-ups you should expect: “Why is greedy safe?” — the argument above; be ready to state it. “What if profits could be negative?” — greedy breaks, since taking one could lock you out later. “What if capital were consumed rather than just required?” — a genuinely different (knapsack-like) problem. “Minimise the number of projects to reach a target?” — same structure, invert the loop condition.

LC 1962 — Remove Stones to Minimize the Total · Medium

Problem. Given pilespiles, where piles[i]piles[i] is a number of stones, you may apply this operation kk times: choose a pile and remove floor(piles[i] / 2)floor(piles[i] / 2) stones from it. Return the minimum total number of stones remaining. (You may pick the same pile more than once.)

Constraints. 1 <= len(piles) <= 10^51 <= len(piles) <= 10^5, 1 <= piles[i] <= 10^41 <= piles[i] <= 10^4, 1 <= k <= 10^51 <= k <= 10^5.

Examples. piles = [5,4,9], k = 2piles = [5,4,9], k = 2 gives 1212 · piles = [4,3,6,7], k = 3piles = [4,3,6,7], k = 3 gives 1212

Editorial — approach, complexity, follow-ups

Removing floor(p / 2)floor(p / 2) from a pile of size pp removes more stones the larger pp is, so each operation should always target the current maximum. Crucially the pile does not disappear — it returns at size p - floor(p/2) = ceil(p/2)p - floor(p/2) = ceil(p/2) and may well still be the largest, so it must go back into the heap.

Time O(n+klogn)O(n + k \log n). Space O(n)O(n).

Three Python details that matter here:

  • heapifyheapify is O(n)O(n), cheaper than nn individual pushes at O(nlogn)O(n \log n).
  • heapreplaceheapreplace is one operation, not a heappopheappop followed by a heappushheappush — it sifts once instead of twice.
  • largest - largest // 2largest - largest // 2 avoids any rounding ambiguity. Writing ceil(largest / 2)ceil(largest / 2) goes through a float, which is unnecessary and imprecise for large values.

The largest <= 1largest <= 1 early exit matters for correctness of the loop’s intent as well as speed: 1 - 1 // 2 == 11 - 1 // 2 == 1, so once every pile is 11, further operations change nothing and the loop would spin kk times pointlessly. ([1,1,1], 5)([1,1,1], 5) returning 33 is that case.

Note this problem needs only one heap — it is included here because it completes the family: the two-heap problems are the ones where a single ordering is insufficient, and seeing where one heap does suffice sharpens the distinction.

Follow-ups you should expect: “Why is greedy optimal?” — the gain from halving is monotonic in pile size, and operations are independent across piles, so an exchange argument shows always taking the largest is at least as good. “What if the operation removed a fixed amount instead?” — greedy still works but piles can empty. “What if you could only use each pile once?” — then sort descending and take the top kk; no heap needed.

LeetCode problem set

#ProblemDifficultyThe twist
1962Remove Stones to Minimize the TotalMediumHalved piles return to the heap; heapifyheapify + heapreplaceheapreplace
2542Maximum Subsequence ScoreMediumSort by the multiplier descending, keep a min-heap of the top kk of the other array
295Find Median from Data StreamHardThe base two-heap balance
502IPOHardUnlock by capital, choose by profit
480Sliding Window MedianHardTwo heaps plus lazy deletion of expired values

Interview follow-ups

They askWhat they’re checkingThe answer
“Why can’t one heap give you the median?”Understanding the structureA heap only orders along root-to-leaf paths; the middle element has no guaranteed position
“How do you get a max-heap in Python?”Language fluencyNegate on push and on read; with tuples, negate only the field you mean to reverse
“Why push-then-shift?”DepthIt maintains the ordering unconditionally — no comparison and no empty-heap guard
“Why let lowlow be the larger half?”PrecisionIt makes the odd-count median unambiguously -low[0]-low[0], with no extra branch
“Sliding window median?”Knowing the hard partHeaps cannot delete arbitrary elements: lazy deletion with a CounterCounter plus explicit size tracking, or SortedListSortedList
“Values are in a small range?”Alternative structuresCounting sort or a Fenwick tree over the value range beats heaps outright
“Why is IPO’s greedy correct?”RigourCapital never decreases, so anything affordable now stays affordable — taking the best available can never reduce future options

Edge-case checklist

  • First insertfindMedianfindMedian after one addNumaddNum must work; push-then- shift handles the empty heap without a guard.
  • Odd vs even counts — one returns a single value, the other a mean; test both, alternating.
  • Return typefloat(-low[0])float(-low[0]) for the odd case, since an intint can fail strict comparison.
  • Duplicate values — perfectly fine; they distribute across both heaps.
  • All identical — the median equals that value; balance still holds.
  • Negative values — easy to break with the negation trick; test them.
  • Nothing affordable (LC 502) — must breakbreak, not loop kk times.
  • kk larger than the input (LC 502, 1962) — the early exit covers it.
  • Piles already all 11 (LC 1962) — halving removes nothing; exit early.

Recap

  • One heap gives you an extreme; two heaps give you the boundary between two halves, which is what a median is.
  • max-heap of the smaller half + min-heap of the larger half. The two tops straddle the median, so reading it is O(1)O(1) and inserting is O(logn)O(\log n).
  • Push-then-shift: push into lowlow, move lowlow’s max into highhigh, then rebalance. Unconditional, no comparisons, no empty-heap case.
  • Let lowlow be the half allowed to be bigger, so the odd-count median is always -low[0]-low[0].
  • Python has no max-heap — negate, and be careful with tuples.
  • The second two-heap shape is unlock by one ordering, choose by another (LC 502): sort by availability, heap by value.
  • Removing an arbitrary element needs lazy deletion (LC 480), or a different structure entirely.

Next: K-way Merge — using a heap to merge many sorted sequences at once.

If this helped you, consider buying me a coffee ☕

Buy me a coffee

Was this page helpful?

Let us know how we did