Median of a stream with heapq

Why it matters

The median is the middle value of a sorted set. For a fixed list, finding it is easy: sort the list and read the middle. The harder question is the one a live system faces. Numbers keep arriving, and after each one you want the median of everything so far. A sorting step on every arrival would cost O(n log n) each time, which adds up quickly when the stream is long.

The median of a stream is a famous interview problem, because it shows how two heaps can cooperate. Each heap is simple on its own. Together they keep the middle value at the boundary between them, with a cost of O(log n) per new number and O(1) to read the answer. The same idea extends to a sliding window, where old numbers also leave, and to a scheduling problem that looks different on the surface but uses the same heap habits.

This lesson builds the two-heap design step by step, traces it on real numbers, and then covers the extensions. You will see the exact output of each program, and you will see where each rule in the code comes from.

The idea

Picture a line of people sorted by height. Split the line down the middle. The shorter half stands on the left, and the taller half stands on the right. The median sits at the split: it is the tallest person on the left when the line has an odd count, or the average of the tallest on the left and the shortest on the right when the count is even.

Now the two halves can be kept in two heaps:

• low is a max-heap holding the smaller half. Its top is the biggest of the small values, which is the value nearest the split from the left. Python's heapq only makes min-heaps, so low stores negated numbers. The smallest negated value is the largest original value.
• high is a min-heap holding the larger half. Its top is the smallest of the big values, which is the value nearest the split from the right.

Two rules keep the structure correct:

1. Every value in low is no larger than every value in high. So the split is clean.
2. The sizes are equal, or low has exactly one more value than high. So the median is always at one of the two tops.

When the count is odd, the median is the top of low, because low holds one extra value. When the count is even, the median is the average of the two tops, because both sit next to the split.

Adding a number, step by step

When a new number arrives, the rules must still hold afterwards. The simplest safe method has three steps.

Step 1: push the number into low. Now low may hold a value that belongs in high, or it may be out of balance in size.

Step 2: move the top of low into high. The largest value in low moves across, so the split is correct again, even if the new number was the largest so far.

Step 3: if high now has more values than low, move the top of high back to low. This fixes rule 2, the sizes, while keeping rule 1, because the smallest value in high is the one that goes back.

Each step is one heap operation, so each new number costs O(log n). Reading the median is O(1), because it only looks at the two tops. The space is O(n), because every number is stored somewhere.

The three steps look odd at first. A value can move to high and then back. The first move makes sure the split is correct even when the new number is the largest so far. The second move fixes the sizes. Think of it as a pipe: every number goes through low, comes out at its top, and lands in the right half. The pipe keeps the order correct without a special case for each situation.

Runnable example: the running median

The first block is the full running-median program. It adds the numbers from a list one at a time. After each number, it prints the median so far. To try other numbers, change the list in the last lines.

The output is:

5.0
10.0
5.0
4.0
5.0
6.0

Read the output line by line. After the first number, 5, the median is 5. After 15, there are two numbers, so the median is the average, (5 + 15) / 2 = 10.0. After 1, there are three numbers, and the middle one is 5. After 3, there are four numbers, and the two middle values are 3 and 5, so the median is 4.0. After 8, there are five numbers, sorted as 1, 3, 5, 8, 15, and the middle is 5. After 7, there are six numbers, sorted as 1, 3, 5, 7, 8, 15, and the middle pair is 5 and 7, giving 6.0.

Now trace the heaps for the first three numbers. Start empty. Add 5: low gets 5, then the top of low (5) moves to high. Now high has one value, larger than low's zero, so the size rule says move it back. low is [5] and high is empty. Median is 5.

Add 15: low gets 15, so it holds 5 and 15. The top of low is 15, which moves to high. Now low is [5] and high is [15]. The sizes are equal, so there is no move back. Median is the average of 5 and 15, which is 10.0.

Add 1: low gets 1, so it holds 1 and 5. The top of low is 5, which moves to high. Now high holds 5 and 15, and low holds 1. High is bigger than low, so the top of high (5) moves back to low. Now low holds 1 and 5, and high holds 15. The median is the top of low, which is 5.

The internal order of a heap is not sorted, so do not expect the list to print in order. Only the tops matter.

Runnable example: watch the two heaps

The next block prints both halves after every number. Printing the halves in sorted order makes the rule easy to check: every value in the left half is at most every value in the right half. The program uses the numbers 2, 4, 6, 8 and 1, and needs no input.

The output is:

after 2: low [2] high []
after 4: low [2] high [4]
after 6: low [4, 2] high [6]
after 8: low [4, 2] high [6, 8]
after 1: low [4, 2, 1] high [6, 8]

Look at the last line. The stream so far is 2, 4, 6, 8, 1. Sorted, it is 1, 2, 4, 6, 8. The five numbers have an odd count, so the median is the middle value, 4. The low side holds the three smallest values, 4, 2, 1 when sorted from largest to smallest, and the top of low is 4. The high side holds 6 and 8. Rule 1 holds, because every value in low is at most 6. Rule 2 holds, because low has one more value than high.

Notice the line after 4. The stream is 2 and 4, so both halves hold one value. low has 2, and high has 4. The median is (2 + 4) / 2, which is 3.0. The printed lists match the split: 2 on the left and 4 on the right.

The visualizer below shows a min-heap, the kind used for the right half. Use it to watch how the smallest value stays at the top while values move in and out. The left half works the same way with the signs flipped, which is how the program above stores it.

Why each design choice matters

The median program has four design choices. Each one has a reason, and interviewers ask about them.

The first choice is storing negated numbers in low. Without the negation, low would be a min-heap, and its top would be the smallest value, which is far from the split. The median needs the largest small value, so low must act as a max-heap. Negation gives that with heapq, which has no max-heap of its own.

The second choice is the three-step add. Without the move from low to high, the largest value stays in low even when it belongs on the right. Rule 1 then breaks, and the top of low can be bigger than the top of high, so the median is wrong.

The third choice is the :.1f in the f-string. It prints one digit after the decimal point. Python always uses a dot as the decimal separator, so the output is the same on every computer.

The fourth choice is the division by 2. In Python 3, / always gives a float, so (1 + 2) / 2 is 1.5 and not 1. The integer division operator // drops the fraction, so it is the one to avoid here. You will see this in the mistakes below.

A full trace of the running median

Use a stream with repeats, because repeats are where mistakes hide. The stream is 3, 1, 4, 1, 5, 9, 2, 6. The table records both heaps after each number. Each side is written in sorted order so it is easy to read, but remember that a heap does not store its values in order. The median column matches the output of the running-median program on this same stream.

Three rows show the three steps of the add at work.

The second row, after 1, shows step 2. The new 1 goes into low, which now holds 1 and 3. The top of low is 3, so it moves to high. Now low holds one value, high holds one, and the sizes are equal, so no move back is needed. The median is the average of the two tops, (1 + 3) / 2, which is 2.0.

The sixth row, after 9, shows the new value going straight across. The 9 is the biggest number so far, so after it goes into low, it is the top of low, and it moves to high at once. The sizes are now three and three. The median is the average of the tops, 3 and 4, which is 3.5.

The seventh row, after 2, shows step 3. The 2 goes into low, and the top of low is now 3. That 3 moves to high, so high has four values and low has three. The size rule is broken, so the top of high, which is the 3 that just moved, moves back to low. Now low holds four values and high holds three. The 3 went across and came back, and the table shows the final state. The median is the top of low, 3.0, because low has one more value.

Notice how little work each row takes. Every add is two or three heap operations, and each operation touches one path in a heap of at most n values. The table has eight rows, and the real program would print eight medians with no sorting at all.

Lazy deletion, step by step

The sliding-window program removes a value every time the window moves, so it needs to delete values from the two heaps. Searching a heap for a value costs O(n), and removing one from the middle breaks the heap order, so the program uses lazy deletion instead. Two ideas work together. The first is a count called delayed, which is a Counter from the collections module. When a value leaves the window, the program adds one to its count there. The second is the cleaning helper prune, which removes values with a count from the top of a heap, but only when they reach the top. A dead value buried deep in a heap stays where it is until it rises.

The program also keeps two integers, low_count and high_count. These count the live values on each side. They are not the heap sizes, because the heap sizes can include dead values that have not yet risen. The median is read only after the tops are cleaned, so a dead value is never read.

Here are the first three slides of the sample input 1 3 -1 -3 5 3 with a window of size 3. The table shows the state after each action. The median column is filled only when the program reads it. Each side lists its live values, top first.

Read the start row first. The window is [1, 3, -1]. Adding 1 puts it in low. Adding 3 is bigger than the top of low, so it goes to high. Adding -1 is at most the top of low, so it goes to low. The counts are two and one. The median is the top of low, which is 1, so the first output is 1.0.

Slide 1 adds -3 and then removes the 1 that just left the window. The add puts -3 into low, which now holds three values. The size rule says low may have at most one more value than high, so the top of low, which is 1, moves to high. Now both sides have two live values. The remove comes next. The value 1 is bigger than the top of low, which is -1, so the program takes the count away from the high side. It records delayed[1] = 1. The top of high is now 1, which is dead, so prune removes it at once. high is left with 3, and the median reads the top of low, -1, so the second output is -1.0.

Slide 2 adds 5 and removes 3. The 5 is bigger than the top of low, so it goes to high. The remove of 3 finds 3 on the high side. The count drops, and the 3 is the top of high, so prune removes it. The median is -1.0, which matches the window [-1, -3, 5], whose middle value is -1.

Slide 3 adds 3 and removes -1. This time the value to remove is on the low side, because -1 is at most the top of low. The count on low drops from two to one, and -1 is the top of low, so prune removes it. Now low has one live value and high has two, so the size rule says high is too big. The top of high, which is 3, moves to low. Now low has two live values, high has one, and low holds -3 and 3, so its top is 3. The answer is 3.0. The window [-3, 5, 3] has middle value 3.

The reason this is correct is simple to state. The counts always match the live values on each side, and the tops are always live after prune runs. The rebalance uses only the counts, so it never moves a dead value. The median reads only the tops, so it never returns a dead value.

An interview conversation

Suppose the interviewer asks for the median of a stream and then asks what changes for a sliding window. A good answer starts with the two-heap design, names the rules, and gives the costs: O(log n) per add and O(1) per median read. Then comes the window. The candidate says that removal from the middle of a heap is the new problem, and that a search would cost O(n). Lazy deletion fixes that, because the value is marked and cleaned only when it reaches the top.

The interviewer may ask whether the heaps can ever disagree about the median. The candidate can answer with the invariant: every live value in low is at most every live value in high, and the live counts differ by at most one. Those two facts are enough to prove the median is always at the tops.

The last question is often about the meeting-room problem, as a check that you can spot the heap pattern in a new setting. A strong answer names the sort by start time and the heap of end times, and it states the rule that the room is free when the end equals the start.

Variations and edge cases

A single value is the first edge case. After one add, low holds it, high is empty, and the median is that value. The code reads the top of low only when low is bigger, so the single value path works without a special case.

Two equal values are the second. Adding 4 and then 4 gives low holding 4 and high holding 4. The median is (4 + 4) / 2, which is 4.0. That is correct, but it shows why the split rule uses "no larger" on the left and "no smaller" on the right. Equal values can sit on either side of the split, as long as the tops are right.

Very large values are the third edge case. In a language with 32-bit integers, the sum of two large values can overflow. Python integers have no size limit, so (2000000000 + 2000000000) / 2 gives 2000000000.0 with no special handling.

Negative values are the fourth. The heaps compare numbers, not sizes, so negatives work the same way. The median of -5, -1 and -3 is -3.0.

A stream with only removals is the fifth, and it cannot happen in the running-median problem, because you need a value to remove first. In the sliding window, the first k adds come before the first remove, so the counts never go negative.

Odd and even counts are the sixth. With an odd count, the median is one value, and the program reads the top of the bigger heap. With an even count, the program averages two tops. Check both when you test. A common off-by-one is to read the top of high when low is bigger. The standard library has statistics.median, which you can use to check your answers on a list. For the stream [2, 4, 6] it gives 4 after all three numbers, which matches the program.

Common mistakes and the messages they produce

Mistake 1: integer division in the median. Writing (low_top + high_top) // 2 gives a whole number. For the pair 3 and 4, (3 + 4) // 2 is 3 rather than 3.5. Use /, which always gives a float in Python 3.

Mistake 2: reading the top of an empty heap. On the first call to median(), if no number has been added, self.low[0] fails:

IndexError: list index out of range

The fix is to add at least one number before reading the median, or to check the sizes first.

Mistake 3: forgetting the size check. Steps 1 and 2 of the add move every new value into high, so high grows and low can be left empty. Try the numbers 1, 2, 3 by hand without step 3. After 1, high holds 1. After 2, high holds 1 and 2. After 3, high holds all three values and low is empty. Reading the median then fails with the same IndexError, because low[0] does not exist.

Mistake 4: using one heap and sorting each time. Each read would cost O(n log n), which is too slow for a long stream. The two-heap design keeps each read at O(1).

Mistake 5: calling list.remove to delete a value from a heap. That search costs O(n), and removing from the middle of a heap list breaks the heap order. It also does not update the counts you track on each side. Use lazy deletion instead, as the sliding-window program does.

Extension: the median of a sliding window

Now a harder version. You have a long list of numbers and a window of size k. The window starts at the first k numbers and moves one step to the right at a time. For each window, print its median.

The running-median design only adds. A window also removes, because the number that leaves the left edge must come out of the structure. Removing a value from the middle of a heap is slow, because the heap must search for it. The trick is lazy deletion. Instead of removing a value at once, you remember that it should be gone. You add a note in a Counter called delayed, counting how many copies of each value need to go. When a value with a note reaches the top of its heap, you remove it then. A value buried deep in the heap keeps its note until it rises.

The program keeps the counts that matter in two integers, low_count and high_count. Those counts tell you how many live values each side holds. The heap sizes can be bigger than the live counts, because they include the values still waiting to be removed. The median is read from the tops, which are always live, because the code removes any dead tops before it reads them.

The rebalance step checks the counts, moves a value from one heap to the other, and then clears dead tops from both heaps. That keeps both rules true at every moment.

Here is the complete program. The window size is 3 and the numbers are built into the code. The program prints one median for each window.

The output is:

1.0
-1.0
-1.0
3.0

The four windows are [1, 3, -1], [3, -1, -3], [-1, -3, 5] and [-3, 5, 3]. Their sorted forms are [-1, 1, 3], [-3, -1, 3], [-3, -1, 5] and [-3, 3, 5]. The medians are the middle values: 1, -1, -1 and 3. Those match the output, which shows the answer is correct.

The enumerate function in the loop gives each number with its position, so nums[i - k] is the number that leaves the window when the position is i.

Trace one step to see the lazy deletion work. When the window moves from [1, 3, -1] to [3, -1, -3], the code first adds -3 and then removes 1. The removal records delayed[1] = 1 and decreases low_count or high_count depending on which side 1 lived on. The value 1 is still physically in its heap until it reaches a top, but the counts no longer count it. When prune runs and 1 is at a top, it is removed. The median is then read from the live tops only.

The method prune is the heart of lazy deletion. It removes dead values only from the top, because only the top can be read. A dead value deep in the heap is harmless until it rises. This keeps each operation cheap. The total cost is O(n log k), because each window add and remove touches heaps of size about k.

The code uses x <= -self.low[0] to decide which side a value lives on. That rule works even when values repeat, because a value equal to the split can sit on either side. The counts stay correct, because each removal decrements exactly one count, and the two heaps keep the values in the right order.

Extension: meeting rooms as a heap problem

The heap also solves a scheduling question that looks different. You have meetings, each with a start time and an end time. A room can be reused as soon as its meeting ends. Find the fewest rooms needed.

The method is greedy. A greedy method makes the best choice it can at each step and does not go back. Sort the meetings by start time. Keep a min-heap of end times, one for each room in use. For each meeting in order, if the earliest end time in the heap is at or before this meeting's start, reuse that room. Pop the heap, then push this meeting's end time. Otherwise, a new room is needed, so just push the end time. At the end, the heap size is the number of rooms.

The rule that a room is free when its meeting ends at the start of the next one is important. A meeting that ends at 30 and another that starts at 30 can share a room, because busy_until[0] <= start is true when both are 30.

The program below uses the four meetings from the example. It needs no input.

The output is:

rooms needed: 2

Trace it. Sorted by start, the meetings are the same order as the input. The first meeting, 0 to 30, pushes 30 onto an empty heap. The second, 5 to 10, looks at the top, 30, which is after 5, so no room is free, and 10 is pushed. The heap is [10, 30]. The third, 15 to 20, looks at the top, 10, which is at or before 15, so the room that ends at 10 is reused. Pop 10, then push 20. The heap is [20, 30]. The fourth, 30 to 40, looks at the top, 20, which is before 30, so reuse again. Pop 20, push 40. The heap is [30, 40], which has size 2. The answer is 2 rooms.

Notice that the heap's size is the answer. The heap only ever holds one entry per room, so its size is the number of rooms needed at once.

Where heaps show up in the same pattern

The meeting-room idea is the same as a general rule: sort the events by time, and keep a heap of the things still in progress. Event simulators use it, because the next event is always the earliest one. Cloud schedulers use it to find the first machine free. A ticket system uses it to find the earliest deadline. Each case asks the same question: which resource frees up first?

Interview patterns from this lesson

The first pattern is two heaps around a split. Use it whenever the question asks for a middle value, a median, or a value at a fixed rank in a stream. Name the two heaps, state the two rules, and show the three-step add.

The second pattern is lazy deletion. Use it when a value must leave a heap, but removing it in the middle would cost O(n). Keep a count of values to delete, and clean the top before you read it.

The third pattern is sort, then a heap of ends. Use it for scheduling and interval problems. Sort by start, and keep a heap of the ends still in use.

Complexity summary

Practice

1. Run the median trace by hand for the input 2, 4, 6. Write low, high and the median after each number. Hint: expect 2.0, 3.0 and 4.0.

Show the answer

After 2: low is [2], high is empty, median 2.0. After 4: low is [2], high is [4], median (2 + 4) / 2 = 3.0. After 6: low is [4, 2], high is [6], and the median is the top of low, which is 4.0.

2. Why must low be a max-heap and not a min-heap? Give the answer in one sentence and one example.

Show the answer

The median needs the largest value on the small side, which sits next to the split, so low must have its largest value on top. For the numbers 1, 2, 3, a min-heap low would have top 1, which is far from the median 2.

3. Show why the integer division (a + b) // 2 gives the wrong median for the pair 3 and 4. What value does Python return, and what should it be?

Show the answer

With the floor division operator, (3 + 4) // 2 is 7 // 2, which is 3. The correct median is 3.5. Use /, which gives 3.5.

4. Change the running-median program so that it ignores a number that is already present. Hint: keep a set of the numbers seen so far, and check it before the add.

Show the answer

Keep seen = set(). At the start of add, write if num in seen: return, then seen.add(num). The in check finds numbers already in the set, so add returns without changing either heap.

5. Remove the third step of the add, the move from high back to low. Predict what happens with the input 1, 2, 3, and say what the median function does next.

Show the answer

Steps 1 and 2 still run. After 1, low is empty and high is [1]. After 2, low is empty and high is [1, 2]. After 3, low is empty and high is [1, 2, 3]. The median function then reads self.low[0], which does not exist, so Python raises IndexError: list index out of range.

6. Explain in your own words why the sliding-window program does not remove a value from the middle of a heap.

Show the answer

Removing a value from the middle of a heap needs a search, which costs O(n), and it breaks the heap order. The program records the removal in a Counter and removes the value only when it reaches a top. Lazy deletion keeps each operation at O(log k).

7. Trace the sliding-window program on the numbers 1 2 3 4 with a window of size 2. Predict the output, then run it to check.

Show the answer

The windows are [1, 2], [2, 3] and [3, 4]. Their medians are 1.5, 2.5 and 3.5. The output is 1.5, 2.5 and 3.5, one per line.

8. Meeting rooms: how many rooms do you need for the meetings (1, 5), (2, 3) and (4, 6)? Show the heap after each meeting.

Show the answer

Sorted by start, the order is the same. (1, 5) pushes 5: heap [5]. (2, 3) looks at 5, which is after 2, so a new room is needed: heap [3, 5]. (4, 6) looks at 3, which is at or before 4, so the room is reused: pop 3, push 6, heap [5, 6]. The answer is 2 rooms.

9. Explain why a meeting that ends at 30 and another that starts at 30 can share a room, and say which comparison in the code allows it.

Show the answer

The room is free at the instant 30. The code uses busy_until[0] <= start, which is true when both values are 30. Changing the comparison to < would force a second room for the back-to-back meetings, so (0, 10) and (10, 20) would need 2 rooms instead of 1.

10. Write a test input for the running-median program with four numbers whose medians are 1.0, 1.5, 2.0 and 2.5. Check it by hand, then run it.

Show the answer

Use the numbers 1, 2, 3 and 4, added in that order. The medians are 1.0 after 1, (1 + 2) / 2 = 1.5 after 2, 2.0 after 3, and (2 + 3) / 2 = 2.5 after 4. The program prints exactly those four lines.

Summary

• The median of a stream uses a max-heap for the smaller half and a min-heap for the larger half, with the median at the two tops. Python stores the max-heap as negated numbers.
• The add rule is: push into low, move the top of low to high, then move the top of high back if high is too big.
• Reading the median is O(1), adding is O(log n), and the space is O(n).
• A sliding-window median adds lazy deletion: keep a Counter of values to remove and clear dead tops before reading.
• Meeting rooms uses a min-heap of end times after sorting by start, and its heap size is the number of rooms.