🎓 A1Learner
🔍
0 of 47 lessons completed

Two Heaps for a Running Median

Two heaps for a running median is keeping a lower half and an upper half of the numbers seen so far. The median is read from the tops of the two heaps after each new number.

Recognition signal

A problem fits this pattern when it asks for the median after each new number arrives. A fresh sort after every number would be slow, so the two halves are kept ready.

signals = [
    "Asks for the median after each new number.",
    "Numbers arrive one at a time.",
    "A fresh sort each time is too slow."
]
for line in signals:
    print(line)

Template: two balanced heaps

Each new number goes to the lower heap, then the largest of the lower half moves up. Rebalancing keeps the sizes within one, so the median is always on top.

import heapq

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

    def add(self, num):
        heapq.heappush(self.low, -num)
        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()
medians = []
for num in [5, 15, 1, 3]:
    finder.add(num)
    medians.append(finder.median())
print(medians)

Worked problem: even counts

With an even count of numbers, the median is the average of the two central values. Those two values sit on top of the two heaps, ready to read.

import heapq

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

    def add(self, num):
        heapq.heappush(self.low, -num)
        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()
medians = []
for num in [2, 4, 6, 8]:
    finder.add(num)
    medians.append(finder.median())
print(medians)

Example

import heapq

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

    def add(self, num):
        heapq.heappush(self.low, -num)
        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()
medians = []
for score in [70, 90, 80, 100]:
    finder.add(score)
    medians.append(finder.median())
print(medians)

Key points

  • Median: the central value of a set of numbers, or the average of the two central values.
  • Lower heap: a heap holding the smaller half, with its largest item on top.
  • Upper heap: a heap holding the larger half, with its smallest item on top.
  • Use negated values in a heap to turn it into a largest-on-top heap.

After 70 the median is 70. After 90 it is the average of 70 and 90. After 80 the central value is 80. After 100 the two central values are 80 and 90, so the median is 85.

Note

the two heaps stay balanced, with sizes differing by at most one.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Which heap holds the smaller half of the numbers?