Median of a stream

Lesson 8.5 · Hard · Two heaps · 15 min · by Vinay C

We'll cover the following

Overview
Watch it run
The idea
Complexity
In Python
Summary

Numbers arrive one at a time. After each one, report the median of everything seen so far.

Watch it run

Press Play or Step to walk through the example, or change the input and press Run. Each frame says what the algorithm does and why.

new410

351

622

974

185

536

708

299

8410

1211

4612

6613

314

9015

41 goes to the large half. Rebalance: move 41 to the small half. The median is 41.

small [41] · large [] · median 41

1 / 16

The idea

Complexity

Time	Space	Why
`O(log n) per add, O(1) per median`	`O(n)`	Each add moves at most three items between the heaps.

In Python

import heapq
class Median:
    def __init__(self): self.lo, self.hi = [], []   # lo is a max-heap (negated)
    def add(self, x):
        heapq.heappush(self.lo, -x)
        heapq.heappush(self.hi, -heapq.heappop(self.lo))
        if len(self.hi) > len(self.lo):
            heapq.heappush(self.lo, -heapq.heappop(self.hi))
    def median(self):
        if len(self.lo) > len(self.hi): return -self.lo[0]
        return (-self.lo[0] + self.hi[0]) / 2

</> Now code it →Try it on LeetCode: Find Median from Data Stream →

Summary: say it in the interview

Split the numbers into a small half (max-heap) and a large half (min-heap), kept within one item of each other. The median is the top of the bigger half, or the average of both tops.
Time O(log n) per add, O(1) per median, space O(n). Each add moves at most three items between the heaps.