Skip to main content

Two Heaps: the Running Median

Editorial

The median sits between two tops

The median of a stream changes with every value, and sorting everything again each time costs O(n log n). But the median only depends on where the two halves meet: the largest value of the lower half and the smallest of the upper half. Keep the lower half in a max-heap and the upper half in a min-heap, with the lower half as big as the upper or one bigger. The median is the lower half's top, or the average of the two tops when the halves are even. An add is O(log n); the median is O(1):

After the stream 5 15 1 3 8 7 2, every value in order on a row of tiles, 1 2 3 5 in orange and 7 8 15 in blue, with 5 in green, split by a dashed line that runs down to a card reading median 5, 4 against 3. Left, a dashed zone titled lower half · max-heap holds 5 over 3 and 2, with 1 under 3, the 5 ringed in green; right, a zone titled upper half · min-heap holds 7 over 15 and 8. Arrows run from each root to the card, labelled its largest and its smallest. A pill reads add O(log n) · median O(1). Focal line: the median is where the two tops meet. Title: the median sits between two tops.

Find the median of a stream

Each new value goes into the lower half; the lower half's largest then moves up to the upper half; and if the upper half is now bigger, its smallest comes back down. No comparison by hand, and the order holds, because only the largest of the lower half ever crosses up and only the smallest of the upper half crosses down:

// The lower half in a max-heap, the upper half in a min-heap; the lower half may hold one more.
class MedianFinder {
    priority_queue<int> low;                                   // its top is the largest of the lower half
    priority_queue<int, vector<int>, greater<int>> high;      // its top is the smallest of the upper half
public:
    void addNum(int num) {
        low.push(num);
        high.push(low.top());                                  // the largest of the lower half moves up
        low.pop();
        if (high.size() > low.size()) {                        // keep the lower half the bigger one
            low.push(high.top());
            high.pop();
        }
    }
    double findMedian() {
        if (low.size() > high.size()) return low.top();
        return ((double)low.top() + high.top()) / 2;           // a double before adding: two ints can overflow
    }
};
import heapq


class MedianFinder:
    """The lower half in a max-heap (as negatives), the upper half in a min-heap; the lower half may hold one more."""

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

    def addNum(self, num):
        heapq.heappush(self.high, -heapq.heappushpop(self.low, -num))   # the largest of the lower half moves up
        if len(self.high) > len(self.low):
            heapq.heappush(self.low, -heapq.heappop(self.high))        # keep the lower half the bigger one

    def findMedian(self):
        if len(self.low) > len(self.high):
            return -self.low[0]
        return (-self.low[0] + self.high[0]) / 2
// The lower half in a max-heap, the upper half in a min-heap; the lower half may hold one more.
static class MedianFinder {
    PriorityQueue<Integer> low = new PriorityQueue<>(Collections.reverseOrder());   // top: the lower half's largest
    PriorityQueue<Integer> high = new PriorityQueue<>();                            // top: the upper half's smallest

    void addNum(int num) {
        low.offer(num);
        high.offer(low.poll());                                 // the largest of the lower half moves up
        if (high.size() > low.size()) low.offer(high.poll());   // keep the lower half the bigger one
    }

    double findMedian() {
        if (low.size() > high.size()) return low.peek();
        return ((double) low.peek() + high.peek()) / 2;         // a double before adding: two ints can overflow
    }
}
// The lower half in a max-heap, the upper half in a min-heap: Introduction to Heaps' MinHeap, twice.
class MedianFinder {
    constructor() {
        this.low = new MinHeap((x, y) => x > y);   // the largest of the lower half on top
        this.high = new MinHeap();                 // the smallest of the upper half on top
    }
    addNum(num) {
        this.low.push(num);
        this.high.push(this.low.pop());            // the largest of the lower half moves up
        if (this.high.a.length > this.low.a.length) this.low.push(this.high.pop());   // the lower half the bigger one
    }
    findMedian() {
        if (this.low.a.length > this.high.a.length) return this.low.a[0];
        return (this.low.a[0] + this.high.a[0]) / 2;
    }
}

Each printed 1.0 1.5 2.0 for Find Median from Data Stream's example, adding 1, 2 and 3, and 5.0 10.0 5.0 4.0 5.0 6.0 5.0 for [5, 15, 1, 3, 8, 7, 2]. Step through the second one:

The running median of [5, 15, 1, 3, 8, 7, 2] — Every value so far, in order: the low half, then the high half, each half's top marked. Green is the median; indigo, the value just added.

Adding 8 when the lower half is 1 3 and the upper half 5 15, in three numbered rows across two dashed zones, the lower half and the upper half, split by a middle line. First, 8 goes into the lower half, now 1 3 8, the 8 in purple. Second, the lower half's largest, 8, crosses the line up into the upper half, now 5 8 15. Third, the upper half has 3 against 2, so its smallest, 5, crosses back: the lower half is 1 3 5 and the upper 8 15. Size pills beside the rows read 3 · 2, 2 · 3 and 3 · 2, and a pill reads 3 against 2: the median is 5. Focal line: only a top ever crosses, so the halves stay in order. Title: only a top crosses.

Below, you keep the halves. Send each value to the low or the high half, then even the sizes by moving a top across. A move that breaks the order or the sizes is refused, and the median shows once the halves are even.

Balance the Halves: place each value, then even the two heaps

Against a sorted list

The simple way keeps every value in a sorted list and reads the middle: insort finds the place in O(log n) but shifts everything after it, O(n) per add. Timed in Python, a median after every add:

ValuesTwo heapsA sorted list
10,00017.2 ms27.7 ms
50,00089.1 ms822.2 ms
100,000397.1 ms2,017.4 ms

At 10,000 values the sorted list is close behind, because its shifts are fast copies in C. Ten times the values cost the heaps 23 times the time, and the sorted list 73 times: O(n log n) in all against O(n²).

A sliding window: deleting late

Sliding Window Median asks for the median of every window of k values. Now values also leave, and a heap can't delete from its middle cheaply. So don't: when a value leaves, mark it and count it out of its half, but leave it in the heap. It's dropped only when it reaches a top, which is the only place it could do harm. Balance by the live counts, and clean a top after every move so both tops are always live:

Three numbered zones showing the same min-heap, the upper half of a window. First, 9 leaves the window: 6 sits over 8 and 7, with 9 and 12 under 8 and 10 under 7, and 9 is drawn pale with a dashed red edge, marked, left in place; 5 live of 6. Second, the values above it go: the heap is now 8 over 9 and 10, with 12 under 9, and 9 still waits; 3 live of 4. Third, 9 is on top, over 12 and 10, and is popped before it's read; then 10 is on top. A legend reads deleted, still in the heap, and live, in the window. Focal line: a deleted value only matters when it's on top. Title: a deleted value waits for the top.
import heapq
from collections import Counter


class DualHeap:
    """A running median that can delete: a deleted value stays in its heap until it reaches the top."""

    def __init__(self):
        self.low, self.high = [], []          # the lower half as negatives (a max-heap), the upper half
        self.gone = Counter()                 # deleted, but maybe still inside a heap
        self.n_low = self.n_high = 0          # live values in each half

    def _clean(self, heap, sign):             # drop deleted values sitting on top
        while heap and self.gone[sign * heap[0]]:
            self.gone[sign * heapq.heappop(heap)] -= 1

    def _balance(self):                       # the lower half holds the extra value, if any
        if self.n_low > self.n_high + 1:
            heapq.heappush(self.high, -heapq.heappop(self.low))
            self.n_low, self.n_high = self.n_low - 1, self.n_high + 1
            self._clean(self.low, -1)
        elif self.n_high > self.n_low:
            heapq.heappush(self.low, -heapq.heappop(self.high))
            self.n_low, self.n_high = self.n_low + 1, self.n_high - 1
            self._clean(self.high, 1)

    def add(self, x):
        if not self.low or x <= -self.low[0]:
            heapq.heappush(self.low, -x)
            self.n_low += 1
        else:
            heapq.heappush(self.high, x)
            self.n_high += 1
        self._balance()

    def remove(self, x):
        self.gone[x] += 1                     # mark it; it leaves the heap when it surfaces
        if x <= -self.low[0]:
            self.n_low -= 1
            if x == -self.low[0]:
                self._clean(self.low, -1)
        else:
            self.n_high -= 1
            if x == self.high[0]:
                self._clean(self.high, 1)
        self._balance()

    def median(self):
        if self.n_low > self.n_high:
            return float(-self.low[0])
        return (-self.low[0] + self.high[0]) / 2


def median_sliding_window(nums, k):
    window = DualHeap()
    for x in nums[:k]:
        window.add(x)
    out = [window.median()]
    for i in range(k, len(nums)):
        window.add(nums[i])
        window.remove(nums[i - k])
        out.append(window.median())
    return out

It returned [1.0, -1.0, -1.0, 3.0, 5.0, 6.0] and [2.0, 3.0, 3.0, 3.0, 2.0, 3.0, 2.0] for the problem's two examples, and agreed with a plain sorted window on 3,000 small cases full of repeated values. Timed on 100,000 values against that sorted window, kept with bisect and insort:

Window kTwo heaps, deleting lateA sorted window
100425.8 ms89.1 ms
10,000376.1 ms473.9 ms
50,000273.4 ms1,162.2 ms

For a small window the sorted window wins: shifting 100 values is one quick copy, while the heaps keep stale values under their tops, 64,135 entries at the peak for a window of 100. As k grows the shifts grow with it, and at k = 50,000 the heaps are 4.3 times faster. In C++, a multiset with an iterator kept on the median does the same job in O(log k) without the marks.

Spotting it, and when it breaks

The clues. "The median" of values that keep arriving, "the middle" of a stream or a window, or any split of a stream into a smaller and a bigger part whose border you need at every step: two heaps facing each other.

The near-misses. One median of a fixed array: quickselect, O(n) on average (Quick Sort and Quickselect). Values from a small range, such as 0 to 100: a count per value, and walk the counts. Medians of short windows in Python: the sorted window, as timed above.

When it breaks. The halves must stay in order and within one of each other's size, after every add. And deleting from the middle of a heap is the expensive part: mark it and drop it late, or use an ordered structure that deletes directly.

Bugs that cost a submission

Integer division (Java).

static class MedianFinder {
    PriorityQueue<Integer> low = new PriorityQueue<>(Collections.reverseOrder());
    PriorityQueue<Integer> high = new PriorityQueue<>();

    void addNum(int num) {
        low.offer(num);
        high.offer(low.poll());
        if (high.size() > low.size()) low.offer(high.poll());
    }

    double findMedian() {
        if (low.size() > high.size()) return low.peek();
        return (low.peek() + high.peek()) / 2;            // int + int, then int / int
    }
}

After adding 1 and 2 it returned 1.0; the answer is 1.5. Two ints added and divided by 2 stay an int, and the half is lost.

Dividing by 2.0 isn't enough (Java).

double findMedian() {
    if (low.size() > high.size()) return low.peek();
    return (low.peek() + high.peek()) / 2.0;          // the int sum overflows first
}

With 2,147,483,647 added twice it returned -1.0: the two ints overflowed when added, before the division saw a double. Sliding Window Median tests exactly this. Cast first, ((double) low.peek() + high.peek()) / 2, as the code above does; it returned 2.147483647E9.

Halves balanced by size alone.

import heapq


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

    def addNum(self, num):
        if len(self.low) <= len(self.high):
            heapq.heappush(self.low, -num)        # the smaller half, whatever the value
        else:
            heapq.heappush(self.high, num)

    def findMedian(self):
        if len(self.low) > len(self.high):
            return -self.low[0]
        return (-self.low[0] + self.high[0]) / 2

Adding 1, 5 and 10 it returned the medians 1.0 3.0 10.0; the answer is 1.0 3.0 5.0. The 10 went into the lower half because it was the smaller half, and became its top. A value has to go to the side its value belongs on, which is what moving the lower half's top across does for free.

Check yourself

Quiz

1 / 6

Multiple choice

Which heap holds the lower half of the values?

Practice problems

In this order:

  1. Find Median from Data Stream: two heaps, the lower half bigger by one at most.
  2. Sliding Window Median: the same, with values that leave deleted late.

Next, Heaps & Priority Queues Quiz: the whole section in sixteen questions.