Streaming Median

Question

Design a class that ingests integers one at a time and can report the median of everything seen so far at any moment. The median is the middle value of the sorted data when the count is odd, and the average of the two middle values when the count is even.

class MedianStream:
    def __init__(self) -> None:
        pass

    def add(self, num: int) -> None:
        pass

    def median(self) -> float:
        pass
Input


 	add(num) receives one integer at a time. median() takes no arguments and is only called after at least one add.

Output


 	add returns nothing. median returns a float, the median of all values added so far.

Examples


 	add(5) then median() returns 5.0 (one value).
 	add(15) then median() returns 10.0 (average of 5 and 15).
 	add(1) then median() returns 5.0 (middle of 1, 5, 15).
 	add(3) then median() returns 4.0 (average of the two middle values 3 and 5 of 1, 3, 5, 15).

Accepted Answer

Maintain two heaps that partition the values around the median: a max-heap lower for the smaller half and a min-heap upper for the larger half. On each insert, push into the max-heap, shift its top across to the min-heap, then rebalance so the max-heap never holds fewer elements than the min-heap. The median is the max-heap top when the total count is odd, or the average of both tops when it is even. Since heapq only gives a min-heap, the max-heap stores negated numbers and negates on the way out. Insert is O(log n), query is O(1), space is O(n).

import heapq

class MedianStream:
    """Online median via two heaps balanced around the middle."""

    def __init__(self) -> None:
        # lower: a max-heap of the smaller half, stored as negatives because
        # heapq is a min-heap (the most-negative value, i.e. the true max,
        # surfaces at index 0).
        self.lower = []
        # upper: a min-heap of the larger half, stored normally.
        self.upper = []

    def add(self, num: int) -> None:
        # Tentatively place num in the lower (max) heap.
        heapq.heappush(self.lower, -num)
        # Move the largest of the lower half into the upper half. This keeps
        # every element of lower  len(self.lower):
            heapq.heappush(self.lower, -heapq.heappop(self.upper))

    def median(self) -> float:
        # Odd count: lower holds one more, so its top is the middle value.
        if len(self.lower) > len(self.upper):
            return float(-self.lower[0])
        # Even count: average the two middle values (the two heap tops).
        return (-self.lower[0] + self.upper[0]) / 2.0

ms = MedianStream()
ms.add(5)
print(ms.median())   # 5.0
ms.add(15)
print(ms.median())   # 10.0
ms.add(1)
print(ms.median())   # 5.0
ms.add(3)
print(ms.median())   # 4.0

 	The invariant after every add is that len(lower) equals len(upper) or is exactly one larger, and every value in lower is at most every value in upper. That is what makes the tops the true middle elements (see the two-heaps note in the Heaps and Priority Queues lesson).

Title	Category	Subcategory	Difficulty	Status
Balanced Brackets	Programming	Other	Easy	Unlock
Circular Queue Implementation	Programming	Other	Hard	Unlock
Evaluate RPN	Programming	Other	Medium	Unlock
Last Node of Singly Linked List	Programming	Other	Medium	Unlock
Linked List Cycle	Programming	Other	Medium	Unlock
Priority Queue Using Heaps	Programming	Other	Hard	Example
Reverse a Linked List	Programming	Other	Easy	Unlock
Validate a BST	Programming	Other	Medium	Unlock
Weather API Data Analyzer	Programming	Other	Hard	Unlock

The Interview Guide for Quants and Traders

Streaming Median

What do you think of this page?