DSA Median from a Data Stream

The problem

A data stream is a sequence of values that arrive one at a time. The median from a data stream is the middle value of everything seen so far. The value must be ready after every single arrival.

An odd count answers the single middle value. An even count answers the average of the two middle values. The script below is the one this article runs.

stream = [5, 15, 1, 3, 8, 7, 9, 2]
expected = [5, 10.0, 5, 4.0, 5, 6.0, 7, 6.0]
print("arrivals:", stream)
print("median after each arrival:", expected)

The median moves both ways as values land. The value 15 pulls it up to 10.0, and the late 2 pulls it back down to 6.0.

Intuition: keep everything sorted

The first idea stores every arrival in a list and keeps the list sorted. The median then sits at the middle index or between the two middle indices. The sorted list below answers each arrival correctly.

seen = []
for value in [5, 15, 1, 3]:
    seen.append(value)
    seen.sort()
    count = len(seen)
    if count % 2 == 1:
        median = seen[count // 2]
    else:
        median = (seen[count // 2 - 1] + seen[count // 2]) / 2
    print("seen sorted:", seen, "median:", median)

The answers are right. The trap is the repeated work. Every arrival re-sorts the whole history, and the history only grows.

Brute force: sort on every query

The brute force is the intuition written as a running process. Each arrival appends one value and sorts everything again. A counter records how many stored values each sort touches.

def running_brute(stream):
    seen = []
    visits = 0
    medians = []
    for value in stream:
        seen.append(value)
        seen.sort()
        visits += len(seen)
        count = len(seen)
        if count % 2 == 1:
            medians.append(seen[count // 2])
        else:
            medians.append((seen[count // 2 - 1] + seen[count // 2]) / 2)
    return medians, visits

medians, visits = running_brute([5, 15, 1, 3, 8, 7, 9, 2])
print("running medians:", medians)
print("values touched by sorting:", visits)

Eight arrivals touch 36 stored values between them. The count is the triangle sum 1 + 2 + ... + 8. A long stream pays that quadratic total in full.

The optimal solution: two heaps split at the middle

Two heaps hold the two halves of the sorted order, the low heap holding the smaller half with its largest value on top and the high heap holding the larger half with its smallest value on top. That one split keeps the median on the heap tops at all times.

A heap in Python is a min-heap, so its smallest entry sits on top. The low side therefore stores negated values, which puts the largest real value on top. Every low value stays at most every high value, and the sizes differ by at most one.

Each arrival enters the low side, sends its top across to the high side, and takes one value back if the high side grew longer. Equal sizes answer the average of the two tops. A longer low side answers its own top.

import heapq


class MedianFinder:
    def __init__(self):
        self.low = []
        self.high = []

    def add(self, value):
        heapq.heappush(self.low, -value)
        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()
stream = [5, 15, 1, 3, 8, 7, 9, 2]
expected = [5, 10.0, 5, 4.0, 5, 6.0, 7, 6.0]
for value, want in zip(stream, expected):
    finder.add(value)
    result = finder.median()
    print("added:", value, "median:", result, "expected:", want, "match:", result == want)

Every arrival matches. No sort ever runs. The median is always a peek at one or two heap tops.

Walkthrough: the first four arrivals

The state is the contents of the two heaps. The low side prints its real values in sorted order, smallest first. Each step adds one value, rebalances, and reads the median off the tops.

import heapq


class MedianFinder:
    def __init__(self):
        self.low = []
        self.high = []

    def add(self, value):
        heapq.heappush(self.low, -value)
        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 value in [5, 15, 1, 3]:
    finder.add(value)
    low_values = sorted(-stored for stored in finder.low)
    high_values = sorted(finder.high)
    print("added:", value, "low holds:", low_values, "high holds:", high_values, "median:", finder.median())

After 15 arrives, the tops are 5 and 15, and their average is 10.0. After 1 arrives, the low side holds 1 and 5, and its top alone is the median. The middle is never searched for, because the split keeps it on top.

Complexity

Each arrival costs a fixed number of heap pushes and pops, and each heap operation costs O(log n) on n stored values. Reading the median costs O(1), a peek at the tops.

The brute force sorts the whole history per arrival, so its total touched values grow with the square of the count. The counts below compare touched values against heap operations at three stream lengths.

import heapq


class CountedFinder:
    def __init__(self):
        self.low = []
        self.high = []
        self.operations = 0

    def add(self, value):
        heapq.heappush(self.low, -value)
        self.operations += 1
        heapq.heappush(self.high, -heapq.heappop(self.low))
        self.operations += 2
        if len(self.high) > len(self.low):
            heapq.heappush(self.low, -heapq.heappop(self.high))
            self.operations += 2

for count in [8, 64, 512]:
    stream = [(i * 37) % 1000 for i in range(count)]
    finder = CountedFinder()
    for value in stream:
        finder.add(value)
    brute_visits = count * (count + 1) // 2
    print("arrivals:", count, "- sorting touches:", brute_visits, " heap operations:", finder.operations)

At 512 arrivals the re-sorting touches 131,328 stored values. The heaps spend at most five operations per arrival. The gap is why streams never re-sort.

Edge cases and pitfalls

A single arrival is its own median. Duplicates straddle the split without any special handling. Negative values work because negation is exact on integers.

import heapq


class MedianFinder:
    def __init__(self):
        self.low = []
        self.high = []

    def add(self, value):
        heapq.heappush(self.low, -value)
        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

def running(stream):
    finder = MedianFinder()
    medians = []
    for value in stream:
        finder.add(value)
        medians.append(finder.median())
    return medians

print("single arrival:", running([42]))
print("two arrivals:", running([1, 2]))
print("duplicates:", running([4, 4, 4]))
print("negatives:", running([-5, -1, -3]))
print("descending:", running([9, 7, 5, 1]))

The classic pitfall is forgetting the negation on the low side. A plain min-heap puts the smallest low value on top, and the median read goes wrong at once.

import heapq

low = []
for value in [5, 15, 1]:
    heapq.heappush(low, value)
print("top of a plain min-heap holding 5, 15, 1:", low[0])
print("the median of 5, 15, 1 needs the largest of the low half: 5")