← Interactive Code Note

Essay

Find the median from a data stream

Use a max heap and a min heap to keep a changing stream balanced and read its median in constant time.

  • Algorithms
  • Go

The goal is to store a growing stream of integers and return its current median in O(1)O(1) time. Sorting the whole collection after every insert would make the read easy, but updates unnecessarily expensive.

Two heaps give us a better boundary.

Split the stream in half

  • A max heap stores the lower half. Its root is the largest value below the median.
  • A min heap stores the upper half. Its root is the smallest value above the median.

Keep their sizes within one element of each other. If both heaps have the same size, the median is the average of their roots. Otherwise it is the root of the larger heap.

Insert and rebalance

  1. Insert into the max heap when the new value is below its root; otherwise insert into the min heap.
  2. When one heap is larger by two elements, move its root to the other heap.
  3. Read one or both roots to calculate the median.

Insertion costs O(logn)O(\log n) because one heap operation may be required. Reading the median remains O(1)O(1).