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^50 <= 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
- Keep
loweras a max-heap (push negated values) andupperas a min-heap. - For each wait: if
loweris empty or the wait is at most the top oflower, push it intolower; otherwise push it intoupper. - If
lowerholds two more thanupper, movelower's root toupper. Ifupperis the larger, moveupper's root tolower. - Report
lower's root whenloweris 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
"""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 boardThe 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, solower[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) // 2on 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.