Running Median with Two Heaps
26 min · Intermediate
Need help? Ask a tutor on WhatsApp (opens in a new tab)The median of a stream changes with every new value. Sorting each time is slow. Two heaps, one for the smaller half and one for the larger half, keep the middle ready at their tops.
The problem and a first answer
A running median is the median of all values seen so far, reported after each new value. For an even count it is the average of the two middle values.
A site records response times one by one and must report the median after each. The median is the middle value of the sorted data. Keeping a sorted list with bisect.insort works.
import bisect
def median_of(sorted_list): mid = len(sorted_list) // 2 if len(sorted_list) % 2: return sorted_list[mid] return (sorted_list[mid - 1] + sorted_list[mid]) / 2
seen = []for x in [5, 2, 8, 1]: bisect.insort(seen, x) print(median_of(seen))Example 1 · Brute force with insort
import bisect
seen = []for x in [7, 3, 9, 1, 5]: bisect.insort(seen, x) mid = len(seen) // 2 m = seen[mid] if len(seen) % 2 else (seen[mid - 1] + seen[mid]) / 2 print(seen, m)[7] 7 [3, 7] 5.0 [3, 7, 9] 7 [1, 3, 7, 9] 5.0 [1, 3, 5, 7, 9] 5
What this code does
- The sorted list shows exactly where the middle is. It is easy to trust, so it makes a good checker for the heap version.
- The cost is O(n) per insert because
insortshifts items.
Reading the median is O(1), but insort shifts items, so each insert costs O(n). For n values that is O(n squared) in the worst case. Sorting again every time is worse, at O(n log n) per step.
Two heaps split the data in half
A two-heap median keeps the lower half in a max-heap and the upper half in a min-heap, sizes equal or low one larger. The median is read from the tops. Each insert costs O(log n).
Keep the smaller half in a max-heap low and the larger half in a min-heap high. The top of each is next to the middle. Keep the sizes equal, or let low hold one extra value.
import heapq
def add(low, high, x): heapq.heappush(low, -x) heapq.heappush(high, -heapq.heappop(low)) if len(high) > len(low): heapq.heappush(low, -heapq.heappop(high))Example 2 · Watch the two heaps
import heapq
def add(low, high, x): heapq.heappush(low, -x) heapq.heappush(high, -heapq.heappop(low)) if len(high) > len(low): heapq.heappush(low, -heapq.heappop(high))
low, high = [], []for x in [7, 3, 9, 1, 5]: add(low, high, x) print(sorted(-v for v in low), sorted(high))[7] [] [3] [7] [3, 7] [9] [1, 3] [7, 9] [1, 3, 5] [7, 9]
What this code does
- The lists show the lower half and the upper half. Every value on the left is at most every value on the right.
- After 5 values
lowholds three andhighholds two, so the median is the top oflow, which is 5.
The new value goes through low first, so its largest value moves across to high. That keeps every value in low at most every value in high. If high is now bigger, move its smallest back.
Reading the median
If low has one extra value, its top is the median. If the sizes are equal, the median is the average of the two tops. Reading takes O(1) time because only the tops are touched.
def median(low, high): if len(low) > len(high): return -low[0] return (-low[0] + high[0]) / 2Example 3 · Median after every value
import heapq
def add(low, high, x): heapq.heappush(low, -x) heapq.heappush(high, -heapq.heappop(low)) if len(high) > len(low): heapq.heappush(low, -heapq.heappop(high))
def median(low, high): if len(low) > len(high): return -low[0] return (-low[0] + high[0]) / 2
low, high = [], []for x in [7, 3, 9, 1, 5, 6]: add(low, high, x) print(median(low, high))7 5.0 7 5.0 5 5.5
What this code does
- An odd count prints an int, such as 7. An even count prints a float, such as 5.0 or 5.5. With six values the two middle ones are 5 and 6, so the last median is 5.5.
- An odd count gives the top of
low. An even count gives the average of the two tops.
Remember low stores negated values, so the real top is -low[0]. The division / gives a float, so an even count can give a half, such as 2.5.
Putting it in a class
Interviews ask for this as a class with two methods: add a number and find the median. Keep both heaps as attributes. The cost is O(log n) to add, O(1) to find, and O(n) space.
class MedianFinder: def __init__(self): self.low, self.high = [], []
def add(self, x): heapq.heappush(self.low, -x) heapq.heappush(self.high, -heapq.heappop(self.low)) if len(self.high) > len(self.low): heapq.heappush(self.low, -heapq.heappop(self.high))
def median(self): if len(self.low) > len(self.high): return -self.low[0] return (-self.low[0] + self.high[0]) / 2Example 4 · A MedianFinder class
import heapq
class MedianFinder: def __init__(self): self.low, self.high = [], []
def add(self, x): heapq.heappush(self.low, -x) heapq.heappush(self.high, -heapq.heappop(self.low)) if len(self.high) > len(self.low): heapq.heappush(self.low, -heapq.heappop(self.high))
def median(self): if len(self.low) > len(self.high): return -self.low[0] return (-self.low[0] + self.high[0]) / 2
finder = MedianFinder()for x in [-4, 10, 2, 2]: finder.add(x) print(finder.median())-4 3.0 2 2.0
What this code does
- Negative numbers and repeats work the same way. The only requirement is that every value in
lowis at most every value inhigh. - Both attributes are plain lists used by
heapq.
Example 5 · Check the heaps against the brute force
import heapqimport bisectimport random
def add(low, high, x): heapq.heappush(low, -x) heapq.heappush(high, -heapq.heappop(low)) if len(high) > len(low): heapq.heappush(low, -heapq.heappop(high))
def median(low, high): if len(low) > len(high): return -low[0] return (-low[0] + high[0]) / 2
random.seed(3)data = [random.randint(-20, 20) for _ in range(200)]low, high, seen, same = [], [], [], Truefor x in data: add(low, high, x) bisect.insort(seen, x) mid = len(seen) // 2 expected = seen[mid] if len(seen) % 2 else (seen[mid - 1] + seen[mid]) / 2 same = same and expected == median(low, high)print(same)True
What this code does
- Both methods agree on every one of 200 random steps. A second, slower solution is a quick way to test a clever one.
- Fixing the random seed makes the test repeat the same data every run.
Call median() only after at least one add, since both heaps are empty at the start. A real version would raise an error or return None for an empty stream.
Next step
Check your understanding
10 quick questions