HEAP-6

Median of a stream

HeapsLesson 6 of 613 min

Values arrive one at a time, and after each one you want the median of everything so far. Two heaps facing each other answer this in O(log n) per value and O(1) per question.

The problem

The median of a list is its middle value once sorted. With an even number of values there are two middle values, and the median is their average. Unlike the average of all values, the median is not pulled far by a few extreme values, which is why it is used for response times, prices and sensor readings.

Simple answers are slow. Sorting after every new value costs O(n log n) each time. Keeping one sorted list and inserting each value in its place costs O(n) per insert, because the values after it must shift. With n values, that is O(n2) in total (n times n).

Two heaps

The median only depends on the values in the middle. So split the values into two halves: low, the smaller half, and high, the larger half. Keep low as a max-heap (its largest value on top) and high as a min-heap (its smallest value on top). The two tops are then the two values nearest the middle.

With that, the median is the top of low when the count is odd, and the average of the two tops when it is even. To add a value x:

  1. If low is empty or x is at most the top of low, push it into low; otherwise push it into high. This keeps the first half of the invariant.
  2. If low now has two more values than high, move the top of low to high. If high now has more than low, move the top of high to low. Moving a top across keeps the order: it is the value closest to the other half.

Step through the figure with the stream 5, 15, 1, 3, 8, 7, 9, 10.

FIG HEAP-6.1The median from two heaps
stream
5
15
1
3
8
7
9
10
lower half: max-heap
upper half: min-heap
median 5
1/8

Add 5. The lower half is empty, so 5 goes there. The count is odd: the median is the top of the lower half.

Step 1 of 8

The code and its cost

import heapq

class MedianFinder:
    def __init__(self):
        # max-heap of the smaller half, as negated values
        self.low = []
        self.high = []  # min-heap of the larger half

    def add(self, x):
        if not self.low or x <= -self.low[0]:
            heapq.heappush(self.low, -x)
        else:
            heapq.heappush(self.high, x)
        if len(self.low) > len(self.high) + 1:
            heapq.heappush(self.high, -heapq.heappop(self.low))
        elif 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 float(-self.low[0])
        return (-self.low[0] + self.high[0]) / 2

Each add does at most three heap operations, so it costs O(log n). Each median reads two tops: O(1). For n values the total is O(n log n), against O(n2) for the sorted list. Practice HEAP-P6 checks this by counting comparisons.

Other orders and windows

  • Any position, not just the middle. To track the 90th percentile, keep low at 90% of the values instead of half. The same two-heap idea works for any fixed fraction.
  • A sliding window. For the median of the last w values only, values must also leave. Use lazy removal (lesson 4): count the values that should leave each heap and drop them when they reach a top. Each step is still O(log w) on average.
  • The k-th smallest of a fixed list. If all the values are there at the start and you need one order position once, quickselect (DVCQ) does it in O(n) on average, with no heap.

HEAP-6 counts as completed once you step through the figure to the last step.