Sliding window median with two heaps in Python

Why it matters

A moving median is a common number in real data. A service might report the median response time of the last 100 requests. A sensor might smooth its readings with the median of the last 10 values. A median is often better than an average for this, because one very slow request does not pull it far away.

The lesson on the median of a stream only adds numbers. Here numbers also leave the window, so the structure has to support removal as well as insertion. That changes the design, and this lesson shows how to handle it with two heaps.

A heap is a list arranged so that its first item is always the smallest. Python provides heaps in the heapq module, which is part of the standard library (the code that comes with Python). You do not need to write one yourself.

By the end of this lesson you will be able to explain why two heaps hold the middle of a set, add and remove numbers while keeping the median easy to read, and check your code against a simple brute-force version.

The problem in plain words

You are given a list of numbers and a window size k. The window is a block of k numbers next to each other. It starts at the left end of the list and slides one step to the right at a time. At each position you need the median of the numbers inside the window.

The median is the middle value when the numbers are sorted. For an odd count, it is the value in the middle. For an even count, it is the average of the two values in the middle.

Here is a small example with k equal to 3. The list is [1, 3, -1, -3, 5].

window [1, 3, -1]     sorted: -1, 1, 3   median: 1
window [3, -1, -3]    sorted: -3, -1, 3  median: -1
window [-1, -3, 5]    sorted: -3, -1, 5  median: -1

The answer for this list is 1, -1, -1. Each window is sorted by hand here, and the middle value is picked. That is the definition. The rest of this lesson is about doing it fast.

The brute-force idea and why it is slow

The simplest method is to copy each window, sort it, and read the middle. In Python that is sorted(nums[i:i + k]). It is correct, and it is a good way to check your faster code. It is slow, though. Sorting k numbers takes about k log k steps, and you do it once per window. For a list of n numbers, the total is about n times k log k steps.

A faster method keeps the window in an order that makes the middle easy to find. Insertion and removal should then be cheap.

The main idea: split the window in half

Split the window's numbers into two groups:

• the lower half, holding the smaller numbers,
• the upper half, holding the larger numbers.

The median lives at the boundary between the two halves. If the window has an odd count, put the extra number in the lower half. Then the median is the largest value in the lower half. If the count is even, the two halves have the same size, and the median is the average of the largest value in the lower half and the smallest value in the upper half.

We need two quick questions answered: the largest value in the lower half, and the smallest value in the upper half. A min heap answers the second, because its first item is its smallest value. heapq makes min heaps only. To get the largest value of the lower half, store each lower-half number as its negative. The smallest stored value is then the largest real value. This is the usual trick in Python for a max heap.

So the design is:

• the lower half is a heap of negated numbers, which acts as a max heap,
• the upper half is a min heap,
• the sizes differ by at most 1, with the lower half allowed to be one larger.

A quick check of the sizes

Take the window [1, 3, -1] from earlier.

• Add 1. It goes to the lower half, which is {1}. The upper half is empty.
• Add 3. The lower half's largest value is 1. Since 3 is larger, it goes to the upper half, which is {3}.
• Add -1. It is smaller than the lower half's largest value, 1, so it goes to the lower half. Lower is {1, -1}, upper is {3}.

The lower half holds two values and the upper half holds one. The window has an odd count, so the median is the largest value in the lower half, which is 1. That matches the sorted list -1, 1, 3.

Sliding the window needs removals

Adding a number is easy with this design. Removing a number is the hard part. A number that leaves the window can be in either heap, and it may be in the middle of its heap, not at the top.

A Python list has a remove method, but it searches the list one item at a time and does not keep the heap order afterwards. Calling it on a heap can break the heap, so you must not do it. The mistakes section shows an example.

The fix is lazy deletion. "Lazy" means the removal is postponed. When a number leaves the window, do not remove it from its heap right away. Instead, write down that it is "owed" a deletion. Then, only when that number reaches the top of its heap, remove it for real.

This works because the median only ever reads the top of each heap. A stale number deep in a heap does no harm. It becomes a problem only when it reaches the top, and at that point we remove it.

Lazy deletion needs two extra things:

• a count, for each heap, of how many deletions are owed for each value,
• separate counters that track how many live numbers each heap holds.

The second point matters. len(heap) includes the stale numbers that are waiting to be removed. The counters give the true size of each half of the window.

For the owed deletions, the code uses Counter from the collections module. A Counter is a dictionary that returns 0 for a value it has never seen, so there is no need to check for a missing key before adding to it.

The worked example, step by step

Use the list [1, 3, -1, -3, 5, 3, 6, 7] with k equal to 3. The first output should be 1.0, the median of [1, 3, -1]. The second should be -1.0, the median of [3, -1, -3].

The code for the walkthrough is in the next section. Here is what happens to the window that starts at index 0 and then the window that starts at index 1.

Window one, from index 0 to 2:

1. Add 1. The lower half is empty, so 1 goes to the lower half. Lower is {1}. Lower size 1, upper size 0.
2. Add 3. The lower half's top is 1. Since 3 is larger, it goes to the upper half. Upper is {3}. Sizes are 1 and 1.
3. Add -1. It is at most 1, so it goes to the lower half. Lower is {1, -1}. Sizes are 2 and 1. The lower half has one more than the upper half, which is allowed.
4. The window is full, so read the median. The lower half is bigger, so the median is its top, 1. The output is 1.0.

Window two, from index 1 to 3. The new number is -3 at index 3, and the number leaving is 1 at index 0.

5. Add -3. The lower top is 1, and -3 is at most 1, so it goes to the lower half. Lower is {1, -1, -3}. Lower size is 3. Upper size is 1. The lower half is now two bigger than the upper half, which breaks the rule.
6. Rebalance. Move the top of the lower half, which is 1, to the upper half. Lower is {-1, -3}, size 2. Upper is {1, 3}, size 2.
7. Remove 1, which is the number leaving the window. The lower half's top is -1. Since 1 is larger than -1, the value 1 is in the upper half. Do not remove it yet. Record one owed deletion for 1 in the upper half's count. The upper half's live size drops to 1.
8. Read the median. Prune the upper half: its top is 1, and 1 has an owed deletion, so remove it now. The upper half is {3}. Lower size is 2 and upper size is 1, so the median is the top of the lower half, which is -1. The output is -1.0.

The steps show the three parts of the method working together: the split, the rebalance, and the lazy deletion.

The Python code

This is the full file. It has the two heaps, the two owed-deletion counts, the live counters, and the sliding loop.

Save the file as sliding_window_median.py and run it:

python3 sliding_window_median.py

Output:

[1.0, -1.0, -1.0, 3.0, 5.0, 6.0]
[1.5, 2.5, 3.5]
[5.0, 5.0, 3.0]

The first line is the example from the walkthrough, and it matches the expected answer for that list. The output prints floats (numbers with a decimal point) such as 1.0 because the median can be an average of two numbers, and Python shows every median as a float.

Reading the code

The add method first calls _prune on the lower heap. The check num <= -self.lower[0] needs a live top, so the stale values must be cleared first. If the lower half has no live values, the number goes to the lower half anyway. That is why the check starts with self.lower_size == 0. Python stops at the first or part that is true, so self.lower[0] is never read on an empty heap.

The remove method uses the same rule to pick a heap. If the number is at most the lower top, it must be in the lower half, because every value in the upper half is at least the lower top. Otherwise it is in the upper half. Equal values can sit in both halves, but they are the same number, so either choice gives the same median.

The median method prunes both tops before reading them. Reading a stale top would give a wrong median. The expression (-self.lower[0] + self.upper[0]) / 2 uses /, which always gives a float in Python 3, even when both numbers are whole.

The _rebalance method keeps the sizes within the rule. When the lower half is two bigger, its top moves to the upper half. When the upper half is bigger, its smallest value moves to the lower half. Both moves prune the source heap first, so the value being moved is a live one. The lower heap stores negated values, so the code negates when it moves a value in or out.

The _prune method is the core of lazy deletion. It looks at the top of a heap. The sign argument is -1 for the lower heap, so sign * heap[0] gives the real value. If that value has an owed deletion, it removes the value and reduces the count. It repeats while the top has owed deletions, because more than one stale value can be stacked at the top. It stops at the first live value. The method is marked @staticmethod because it does not use self.

The main loop adds the new number first and removes the old one second. Exercise 3 shows that the order does not change the output.

The second example: repeated values

The third line of the output comes from the list [5, 5, 5, 1] with k equal to 2. The windows are [5, 5], [5, 5] and [5, 1].

• Window [5, 5]: both values are 5, so the median is 5.0.
• Window [5, 5] again after the next step: still 5.0.
• Window [5, 1]: sorted is 1, 5, and the median is their average, 3.0.

This example shows why the deletion counts are kept in a Counter, not a set. When one 5 leaves the window, the other 5 is still in it. A set would forget the value 5 completely, and the window would lose a number. A count keeps track of how many copies are owed a deletion.

Checking against brute force

Hand-checking every case is slow. The lesson's code was checked against the simple version from the brute-force idea: copy each window, sort it, and read the middle.

The check used 5,000 random lists, generated with a fixed seed, so the run can be repeated. Each list had 1 to 30 numbers taken from the small range -3 to 3, so duplicates were common. For every window in every list, the two answers were compared.

total mismatches: 0

Why this design is fast

Each number is added once and removed once. Each add or remove touches one heap, and each heap operation takes about log of the heap size steps, which is close to log k. Reading the median takes a constant number of steps after pruning. Each number is also pruned at most once, because a stale value is removed from its heap when it reaches the top and never returns.

So the whole list of n numbers takes about n log k steps, instead of the n times k log k steps of the brute-force method.

The cost of lazy deletion is memory. The heaps can hold stale values until they reach the top. In the worst case the heaps hold up to about n numbers, not just k. For a window of a few thousand numbers this is usually fine.

Common mistakes

Using the heap length as the window size

len(self.lower) counts stale values too. A value that left the window but has not yet reached the top still counts. If the code uses len(self.lower) instead of self.lower_size, the rebalance rule runs on the wrong numbers. The median can then come from the wrong half, and the output looks plausible but is wrong. The fix is to keep the separate live counters.

Forgetting to prune before reading a top

A stale value at the top of a heap is a number that is no longer in the window. If the median reads it without pruning, the output contains a value that should be gone. The median, add, remove and _rebalance methods each prune before they read a top, and that rule should hold in any change you make.

Forgetting the negative sign for the lower half

heapq puts the smallest item first. If the lower half stores the numbers as they are, its first item is the smallest number, not the largest. Here is a small check:

Output:

heapq top after pushing 4, 1, 7: 1

The top is 1, not 7. The lower half needs the negated numbers, as the code above does.

Using // for the average

// is floor division. It rounds down, so for negative numbers the result is further from the true average than you might expect. Here is the difference:

Output:

floor division of -3 and -2: -3
true division of -3 and -2: -2.5

The median of a window with values -3 and -2 is -2.5. The code uses / for that reason.

Calling list.remove on a heap

A list's remove method finds a value by searching the list. It does not keep the heap order. Here is a heap that loses its order after remove:

Output:

heap list: [2, 4, 10, 18, 6, 17]
after list.remove(2): [4, 10, 18, 6, 17]
popped in order: [4, 10, 6, 17, 18]

The sorted order should be 4, 6, 10, 17, 18. The output has 10 before 6. Removing an item from the middle of the list broke the order that heapq relies on, and heapq does not repair it. Lazy deletion avoids this problem.

Reading the median of an empty window

If the window is empty and the code asks for the median, self.lower[0] looks at an empty list. Here is the message:

IndexError: list index out of range

The median is only read when the window has at least one value. The main loop reads the median after both the add and the remove, so the window always has a value at that point.

Python does not overflow, but / returns a float

Python integers have no fixed size. 2**62 + 2**62 gives 9223372036854775808 in Python, with no wrap-around to a negative number as can happen with a fixed-size integer in some other languages. So the average of two large numbers does not overflow.

Floats have less room. A float keeps about 15 to 16 decimal digits, so a very large average can lose its last digits. In this run, (10**20 + 1) / 2 printed 5e+19, and the + 1 was lost. (10**20 + 1) // 2 gives the exact result, 50000000000000000000. The median function uses / because its inputs are normal-sized numbers in this lesson. For very large values, use // when the exact result is needed and both values are whole.

Practice

Each exercise has its answer below it. Work the answer out on paper before you read it.

Exercise 1: trace a small window

Find the output of the sliding window median for [2, 4, 6] with k equal to 2.

Show the answer

The windows are [2, 4] and [4, 6]. The medians are the averages of each pair, 3.0 and 5.0. The output is [3.0, 5.0]. The code prints the same values, which you can check by running the file with that input.

Exercise 2: repeated values

Predict the output for [1, 1, 1] with k equal to 3. Then explain what would be lost if the deletion counts were kept in a set instead of a Counter, using the window [5, 5, 1] as the example.

Show the answer

The output is [1.0]. There is one window, and all three values are 1. For the second part, a set would contain 5 once. When one 5 leaves the window, the set would remove 5 completely, even though another 5 is still in the window. The median would then be computed from a list that is missing a value. A count of 2 for the value 5 keeps the second copy in place until it also leaves.

Exercise 3: what happens if removal runs first

The main loop adds the new number first and then removes the old one. Explain what happens with k equal to 1 if the order were reversed.

Show the answer

Nothing goes wrong. With k equal to 1, the window holds one number before the new one arrives, and that number is in the lower half. Removing it first sends the removal to the lower half, the window is empty for a moment, and then the new number is added to the lower half again. The median is only read after the add, so the output is [4.0, 2.0, 8.0] in both orders for the input [4, 2, 8]. A check of 5,000 random lists, with every window of each list run in the reversed order, found no difference from the brute-force answers.

Exercise 4: lower median for even windows

Change the method so that for an even window it returns the lower of the two middle values, instead of their average. Write the new median method.

Show the answer

Because the lower half holds the extra number when the window size is odd, and equal halves when it is even, the largest value of the lower half is the lower middle value in both cases. After pruning, self.lower[0] gives that value, so the method only needs to return it:

def median(self):
    self._prune(self.lower, self.lower_delayed, -1)
    return float(-self.lower[0])

The method still returns a float, so the caller does not change. The top of the lower heap is always live after pruning, because the window has at least one value. A check of 3,000 random lists against a sorted copy of each window found no mismatches for this version. For the input [1, 2, 3, 4] with k equal to 2, it prints [1.0, 2.0, 3.0].

Exercise 5: a stale value at the top

Suppose the call that prunes the lower heap were removed from median, and the lower heap's top was a value that had an owed deletion. Explain why the median could be wrong, and give a small list that would show it.

Show the answer

The median would read the stale value as if it were still in the window. Take the list [-1, -2, -1] with k equal to 2. The windows are [-1, -2] and [-2, -1]. The correct medians are -1.5 and -1.5.

When the second window is formed, the first number, -1, leaves. The stale value -1 sits in the lower half with an owed deletion, and it is the lower top, because the lower half's largest value is -1. Without pruning, the median would be (-1 + -1) / 2, which is -1.0. With the _prune call in median, the stale -1 is removed first, the lower top becomes -2, and the median is (-2 + -1) / 2, which is -1.5. A run of the code with median changed in this way printed [-1.5, -1.0] where the correct output is [-1.5, -1.5].

Summary

• The median of a window sits at the boundary of a max heap for the smaller half and a min heap for the larger half.
• Python's heapq makes min heaps only, so the lower half stores negated numbers.
• The two halves stay within one value of each other in size, and the lower half may hold the extra value.
• Lazy deletion records owed removals in a Counter and removes a value only when it reaches the top of its heap.
• Keep separate live counters, because the heap length includes stale values.
• Prune before reading a heap top, and use / for an average so the result is a float.