Find Median from Data Stream

Q: What is the time and space complexity of Find Median from Data Stream?

Time O(log n) add, O(1) median, space O(n).

Hard

Balance a max-heap of the low half against a min-heap of the high half.

Two Heaps, Number by Number

low keeps the smaller half with its largest on top; high keeps the larger half with its smallest on top. Both start empty.

StopHeapExport

stream

CH1 low
CH2 high

Problem

Build a data structure that accepts integers incrementally and reports the median of the values currently stored. The median is the middle value when the numbers are sorted, or the average of the two middle values when the count is even. Support insertion and median queries efficiently.

Examples

Input: addNum(5), addNum(15), findMedian(), addNum(1), findMedian(), addNum(3), findMedian()

Output: 10.0, 5, 4.0

After [5, 15], the median is (5 + 15) / 2 = 10. After [1, 5, 15], it is 5. After [1, 3, 5, 15], it is (3 + 5) / 2 = 4.

Input: addNum(2), addNum(8), addNum(4), findMedian()

Output: 4

The sorted values are [2, 4, 8], so the middle value is 4.

The Goal

Numbers arrive one at a time, and you may be asked for the median after each: the middle number, or the average of the middle two when the count is even.

For the stream 5, 15, 1, 3, the medians are 5, 10, 5 and 4.

Sorting after every number repeats work. Aim for O(log n) per number and O(1) per median.

Keep the middle in reach

Split the numbers into a smaller and a larger half. A max-heap shows the top of the smaller half; a min-heap, the bottom of the larger.

Keep the halves balanced, and never sort.

Key Idea: Two Heaps

low is a max-heap holding the smaller half. high is a min-heap holding the larger half. Every number in low is at most every number in high.

Keep low the same size as high, or one bigger. Then an odd count's median is the root of low, and an even count's is the average of the two roots.

Everything in low is at most everything in high.

low: 3, 1. high: 5, 15.

The Algorithm

python

import heapq


class MedianFinder:
    def __init__(self):
        # low: the smaller half, stored negated so its root is its largest.
        # high: the larger half, its smallest at the root.
        self.low = []
        self.high = []

    def addNum(self, num):
        if not self.low or num <= -self.low[0]:
            heapq.heappush(self.low, -num)
        else:
            heapq.heappush(self.high, num)
        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 findMedian(self):
        if len(self.low) > len(self.high):
            return -self.low[0]
        return (-self.low[0] + self.high[0]) / 2

Python's heapq only builds min-heaps, so low stores each number negated. Its smallest stored number is its largest real one.

A new number goes in low if it's at most the largest in low, and in high otherwise. Then one root moves across, if the sizes need it.

Approach	addNum	findMedian	Trade-off
Sort after every number	O(n log n)	O(1)	Repeats the ordering
Unsorted list	O(1)	O(n)	Every median scans
Two heaps	O(log n)	O(1)	Keeps only the middle in reach

Walkthrough

Run it on 5, 15, 1, 3.

5 goes in low and is the median. 15 is more than 5, so it goes in high, and the median is (5 + 15) / 2 = 10. 1 goes in low, which then holds the extra number: the median is 5.

3 goes in low too, making it two bigger, so 5 moves across. low holds 3 and 1, high 5 and 15, and the median is (3 + 5) / 2 = 4.

Checkpoint · predict the next step

Step 5 of 14 on the stage: Add 15.

low = [5] · high = [] · median 5

What does step 6 show?low = [5, 1] · high = [15] · median 10low = [5] · high = [15] · median 5low = [5] · high = [] · median 5low = [] · high = []

Step through it on the stage

Pitfalls

Rebalance after every number, even when no median is asked for. If low held 1, 2 and 100 and high held 101, findMedian would return 100. The median of 1, 2, 100, 101 is 51.
Move the right root: the largest from low, the smallest from high.
Use true division. -3 and 4 give 0.5.
Equal numbers need nothing special. Try Break it with five 4s.

After 1, 2, 100 and 101, which state is right?

Complexity

Each addNum does one push, plus at most one pop and push to rebalance, each O(log n). findMedian reads one or two roots, so it's O(1).

Every number stays in one of the heaps, so space is O(n).

Two heaps

Time

O(log n)

Space

O(n)

Where the time goes

Push a numberA heap's height is log n.dominantO(log n)
RebalanceAt most one root moves across.dominantO(log n)
Read the medianOne root, or two.O(1)
Store every numberEach sits in one heap.O(n)

How it grows

Roughly what that costs

n	10	100	1,000	10,000
O(log n)	3	7	10	13
O(n)	10	100	1K	10K

Logarithmic adds, constant-time medians, linear space.

In short

Keep a max-heap for the smaller half and a min-heap for the larger half, rebalancing so their sizes differ by at most one. The median is the top of the larger heap (odd count) or the average of both tops (even). Each insert is O(log n); median is O(1).

TimeO(log n) add, O(1) medianSpaceO(n)

Common mistake

Letting the heaps drift out of balance, so the tops no longer straddle the true median.

NextKth Largest Element in an Array →A size-k min-heap keeps exactly the k largest seen so far.

Kth Largest Element in an ArrayMedium Top K Frequent ElementsMedium