HeapshardTwo heaps balanced around the middle4 min · 127 of 290

The board shows the median

Report the median wait after every patient seen, by keeping the lower half under a max-heap and the upper half under a min-heap.

An urgent-care clinic logs each patient's wait as they are called through, and the board updates. Re-sorting the day after every patient works, and is quadratic.

The problem

Reception records how many whole minutes each patient waited, at the moment they are called through. The board then shows the median wait of everyone seen so far today: the middle value of the day's waits in order, or the average of the two middle values when an even number of patients has been seen.

The median is posted rather than the average because one four-hour wait should not make the other forty people look like they waited an hour. The board updates after every patient, so the answer is a running one — a value per patient, not a single number at the end.

Input. waits — a list of integers, the wait in minutes of each patient, in the order they were seen.

Output. A list of the same length. Entry i is the median of the first i + 1 waits, as a floating-point number.

Example.

waits = [12, 4, 9]   ->  [12.0, 8.0, 9.0]

One patient: the median is that wait. Two: the average of 4 and 12. Three: the middle of 4, 9 and 12.

A second example, where long waits keep landing on the high side:

waits = [2, 40, 3, 41, 4]   ->  [2.0, 21.0, 3.0, 21.5, 4.0]

After four patients the ordered waits are 2, 3, 40, 41 and the board reads 21.5 — the two middle values averaged, not rounded.

Constraints.

  • 0 <= len(waits) <= 2·10^5
  • 0 <= waits[i] <= 10^4
  • an empty list of patients produces an empty board

Hints

Hint 1

Re-sorting the day after every patient is correct. How much of that sorted run do you read?

Hint 2

Split the day in two halves at the middle. What single value do you need from each half, and where does a heap keep exactly that value?

Hint 3

A new wait can land on either side. Two things must be restored after each insertion: which half a value belongs in, and how big the halves are.

Approach

Brute force

Keep the waits in a list, sort after each patient, read the middle: n sorts of a growing run, so about n² log n comparisons — around 10¹¹ for 2·10⁵ patients. Inserting into a sorted list with bisect finds the spot in log n but still shifts the tail, so it stays quadratic: 4·10¹⁰ element moves.

The insight

The board never needs the day sorted, only its two middle values — so keep a max-heap of the lower half and a min-heap of the upper half and read their two roots.

A heap gives the largest of the lower half and the smallest of the upper half in constant time. Those two roots straddle the middle exactly when two conditions hold: every value in the lower half is at most every value in the upper half, and the halves differ in size by at most one. Both are restorable in log n after each insertion by moving one root across.

Algorithm

  1. Keep lower as a max-heap (push negated values) and upper as a min-heap.
  2. For each wait: if lower is empty or the wait is at most the top of lower, push it into lower; otherwise push it into upper.
  3. If lower holds two more than upper, move lower's root to upper. If upper is the larger, move upper's root to lower.
  4. Report lower's root when lower is larger, otherwise the average of the two roots.

Complexity

Time O(n log n) — one push, at most one transfer and a constant-time read per patient. Space O(n): every wait sits in one of the two heaps.

Solution

Python 3 · standard library30 lines · 7 test cases, all passing
"""The board shows the median — two heaps kept balanced around the middle wait."""

import heapq


def solve(waits):
    lower = []       # a max-heap, stored negated: the smaller half of the waits
    upper = []       # a min-heap: the larger half
    board = []

    for wait in waits:
        # Invariant: everything in lower is <= everything in upper, and lower holds
        # either the same number of waits as upper or exactly one more. The middle
        # of the run is therefore always at one of the two roots.
        if not lower or wait <= -lower[0]:
            heapq.heappush(lower, -wait)
        else:
            heapq.heappush(upper, wait)

        if len(lower) > len(upper) + 1:
            heapq.heappush(upper, -heapq.heappop(lower))
        elif len(upper) > len(lower):
            heapq.heappush(lower, -heapq.heappop(upper))

        if len(lower) > len(upper):
            board.append(float(-lower[0]))
        else:
            board.append((-lower[0] + upper[0]) / 2)

    return board
The cases that ran
TESTS = [
    (([12, 4, 9],), [12.0, 8.0, 9.0]),
    (([9, 7, 5, 3],), [9.0, 8.0, 7.0, 6.0]),
    (([1, 2],), [1.0, 1.5]),
    (([5, 5, 5, 5],), [5.0, 5.0, 5.0, 5.0]),
    (([7],), [7.0]),
    (([],), []),
    (([2, 40, 3, 41, 4],), [2.0, 21.0, 3.0, 21.5, 4.0]),
]

Pitfalls

  • Comparing against lower[0] without negating it. The lower half is stored negated, so lower[0] is minus the largest value in it. Routing on the raw root sends waits to the wrong heap, and [12, 4, 9] prints 12.0, 8.0, 12.0 — the third patient reads a median no one waited near.
  • Averaging with integer division. (a + b) // 2 on the second example gives 21 rather than 21.5, and the board loses half a minute on every even-numbered patient.
  • Reading lower's root whatever the sizes are. With equal halves that is the lower of the two middle values, so [1, 2] posts 1.0, 1.0.

Variants

  • The k lightest crates — one heap drained once, where this keeps two of them balanced all day.
  • Top-k and two heaps — the lesson that sets out the balance rule and what breaks when it slips by two.