Finding the Median of a Stream

Split the numbers into two halves with two heaps: a max heap holding the smaller half and a min heap holding the larger half. Keep the sizes balanced so that the max heap holds either the same number of elements as the min heap or exactly one more.

With that invariant, the median is the root of the max heap when the total count is odd, and the average of the two roots when it is even. Each insertion is one push into a heap plus at most one move across the boundary, so it is logarithmic, and the query itself is constant time.

import heapq


class MedianFinder:
    def __init__(self):
        self.lower = []
        self.upper = []

    def add(self, value):
        if not self.lower or value <= -self.lower[0]:
            heapq.heappush(self.lower, -value)
        else:
            heapq.heappush(self.upper, value)
        if len(self.lower) > len(self.upper) + 1:
            heapq.heappush(self.upper, -heapq.heappop(self.lower))
        else:
            if len(self.upper) > len(self.lower):
                heapq.heappush(self.lower, -heapq.heappop(self.upper))

    def median(self):
        if len(self.lower) > len(self.upper):
            return -self.lower[0]
        return (-self.lower[0] + self.upper[0]) / 2


finder = MedianFinder()
for value in [5, 15, 1, 3]:
    finder.add(value)
print(finder.median())

The output is 4.0. The lower heap holds 1 and 3 and the upper heap holds 5 and 15, so with four values the median is the average of the two middle ones.

The mechanism is that the ordering between the two heaps is maintained by the insertion rule. A new value is compared with the root of the lower heap, which is the largest of the smaller half, and it goes to whichever side it belongs on. The rebalance then moves one element across so that the sizes stay right, and because only one element ever moves, the ordering between the halves is preserved.

The trade-off to name is the constant factors and the memory, since every value is stored in one of the heaps and the negation doubles the arithmetic on the lower side. It also matters that you decide what to return for an empty stream, because the median does not exist then. The advantage over keeping all the values in a sorted list is that you never pay to insert into the middle, so the cost per element is logarithmic instead of linear.

Keep track of which heap each value went into, and when removing, take it out of that heap and then restore the size balance by moving an element across the boundary if needed.

import heapq


class MedianFinder:
    def __init__(self):
        self.lower = []
        self.upper = []

    def add(self, value):
        if not self.lower or value <= -self.lower[0]:
            heapq.heappush(self.lower, -value)
        else:
            heapq.heappush(self.upper, value)
        self._rebalance()

    def _rebalance(self):
        if len(self.lower) > len(self.upper) + 1:
            heapq.heappush(self.upper, -heapq.heappop(self.lower))
        elif len(self.upper) > len(self.lower):
            heapq.heappush(self.lower, -heapq.heappop(self.upper))

    def remove(self, value):
        if value <= -self.lower[0]:
            self.lower.remove(-value)
            heapq.heapify(self.lower)
        else:
            self.upper.remove(value)
            heapq.heapify(self.upper)
        self._rebalance()

    def median(self):
        if not self.lower and not self.upper:
            return None
        if len(self.lower) > len(self.upper):
            return -self.lower[0]
        return (-self.lower[0] + self.upper[0]) / 2


finder = MedianFinder()
for value in [5, 15, 1, 3]:
    finder.add(value)
print(finder.median())
finder.remove(1)
print(finder.median())

The outputs are 4.0 and 3.0. After removing the 1 the remaining values are 3, 5 and 15, and the median of those three is 3.

The mechanism is that removing from a heap is not a constant time operation, so the list has to be repaired with a fresh heapify. The lower heap stores negated values, which is why the removal searches for the negated number. The rebalance step then restores the size invariant, and because a single element moved across the boundary, the ordering between the heaps is still correct.

The trade-off to name is that removal is linear in the size of the heap because of the linear scan and the reheapify, so this structure is only appropriate when removals are rare or when the removed element is known to be at a root. If you need frequent arbitrary removals, the answer is to keep the values in a sorted structure such as a balanced tree with order statistics, or to use a lazy deletion heap where removed elements are marked and discarded when they surface.