Solving Find Median from Data Stream in Python: Step-by-Step Guide

Introduction to Find Median from Data Stream

The "Find Median from Data Stream" problem is a classic algorithmic challenge frequently encountered in coding interviews and real-world data processing scenarios. The task is deceptively simple: design a data structure that supports adding numbers from a continuous stream and efficiently retrieving the median of all numbers seen so far at any point in time.

While computing the median of a static list is straightforward—sort the list and pick the middle element—doing so repeatedly as new data arrives presents a unique challenge. Naively re-sorting the entire collection after every insertion leads to poor performance, especially when dealing with high-frequency data streams such as stock prices, sensor readings, or network traffic metrics.

What Is the Median?

The median is the middle value in an ordered dataset. If the dataset has an odd number of elements, the median is the central element. If the count is even, the median is typically the average of the two central elements. For example, in the sorted list [1, 3, 5], the median is 3. In [1, 3, 5, 7], the median is (3 + 5) / 2 = 4.0.

Why This Problem Matters

Streaming median computation has significant practical applications across multiple domains:

Financial analytics: Tracking the median stock price or trade volume in real time helps identify anomalies and trends without being skewed by outliers.
Monitoring systems: Computing median latency for API requests provides a more robust performance metric than the mean, which is sensitive to extreme values.
IoT and sensor networks: Aggregating median readings from distributed sensors supports reliable environmental monitoring.
Machine learning: Online median calculations are useful for feature normalization and outlier detection in streaming data pipelines.

From an interview perspective, this problem elegantly tests a candidate's understanding of heap data structures, algorithmic trade-offs, and object-oriented design. Mastering it demonstrates the ability to move beyond brute-force solutions toward optimal, scalable designs.

Understanding the Naive Approach

Before diving into the optimal solution, it is instructive to examine the naive approach. The simplest implementation stores all numbers in a list, sorts it whenever the median is requested, and returns the middle element(s).

class NaiveMedianFinder:
    def __init__(self):
        self.data = []

    def add_num(self, num):
        self.data.append(num)

    def find_median(self):
        self.data.sort()
        n = len(self.data)
        if n % 2 == 1:
            return self.data[n // 2]
        else:
            return (self.data[n // 2 - 1] + self.data[n // 2]) / 2

While correct, this approach has serious performance drawbacks. Each call to find_median triggers a full sort, costing O(n log n) time. If the median is queried frequently—say, after every insertion—the overall cost becomes O(n^2 log n), which is unacceptable for large streams.

The Optimal Solution: Two Heaps

The key insight for an efficient solution is to maintain the dataset split into two halves at all times: a lower half containing the smaller numbers and an upper half containing the larger numbers. By using two heaps, we can access the largest element of the lower half and the smallest element of the upper half in constant time, which gives us direct access to the median.

How the Two-Heap Strategy Works

We use two heaps:

Max-heap (lower half): Stores the smaller half of the numbers. The largest element of this half sits at the top, giving us quick access to the left side of the median.
Min-heap (upper half): Stores the larger half of the numbers. The smallest element of this half sits at the top, giving us quick access to the right side of the median.

Python's heapq module provides a min-heap implementation. To simulate a max-heap, we store negated values. This trick allows us to reuse the same efficient heap operations while logically treating the structure as a max-heap.

The algorithm maintains two invariants:

All elements in the max-heap are less than or equal to all elements in the min-heap.
The sizes of the two heaps differ by at most one. If the total count is odd, the max-heap holds the extra element, making its top the median.

Step-by-Step Insertion Logic

When a new number arrives, we follow a careful insertion procedure to preserve the invariants:

If the max-heap is empty or the new number is less than or equal to the top of the max-heap, push it onto the max-heap (as a negated value).
Otherwise, push it onto the min-heap.
Rebalance the heaps: if one heap has more than one element than the other, move the top element from the larger heap to the smaller one.

This rebalancing step ensures the size invariant always holds, keeping the median accessible in O(1) time.

Complete Python Implementation

Below is the complete, production-ready implementation of the two-heap median finder:

import heapq


class MedianFinder:
    def __init__(self):
        # Max-heap for the lower half (store negated values)
        self.lower_half = []
        # Min-heap for the upper half
        self.upper_half = []

    def add_num(self, num):
        # Step 1: Decide which heap to insert into
        if not self.lower_half or num <= -self.lower_half[0]:
            heapq.heappush(self.lower_half, -num)
        else:
            heapq.heappush(self.upper_half, num)

        # Step 2: Rebalance the heaps
        if len(self.lower_half) > len(self.upper_half) + 1:
            moved = -heapq.heappop(self.lower_half)
            heapq.heappush(self.upper_half, moved)
        elif len(self.upper_half) > len(self.lower_half):
            moved = heapq.heappop(self.upper_half)
            heapq.heappush(self.lower_half, -moved)

    def find_median(self):
        if len(self.lower_half) > len(self.upper_half):
            return float(-self.lower_half[0])
        else:
            return (-self.lower_half[0] + self.upper_half[0]) / 2.0

Walking Through an Example

Let us trace through a sequence of insertions to see how the data structure evolves. Suppose we add the numbers 5, 2, 8, 1, 9 one at a time:

finder = MedianFinder()

finder.add_num(5)
# lower_half: [-5], upper_half: []
# median = 5.0

finder.add_num(2)
# 2 <= 5, push to lower_half, then rebalance
# lower_half: [-2], upper_half: [5]
# median = (2 + 5) / 2 = 3.5

finder.add_num(8)
# 8 > 2, push to upper_half
# lower_half: [-2], upper_half: [5, 8]
# median = (2 + 5) / 2 = 3.5

finder.add_num(1)
# 1 <= 2, push to lower_half, then rebalance
# lower_half: [-2, -1], upper_half: [5, 8]
# median = 2.0

finder.add_num(9)
# 9 > 2, push to upper_half
# lower_half: [-2, -1], upper_half: [5, 8, 9]
# rebalance moves 5 to lower_half
# lower_half: [-5, -1, -2], upper_half: [8, 9]
# median = 5.0

At each step, the median is computed in constant time by inspecting only the tops of the heaps. The sorted order of the full dataset [1, 2, 5, 8, 9] confirms the final median of 5.0.

Complexity Analysis

The two-heap approach offers excellent performance characteristics:

add_num: O(log n) time due to heap push and pop operations. Space is O(n) overall to store all elements.
find_median: O(1) time since we only read the top elements of the heaps.

Compared to the naive O(n log n) per median query, this represents a dramatic improvement, especially when median queries are frequent relative to insertions. For a stream of one million numbers with a median query after each insertion, the two-heap solution performs roughly one million times faster than the naive approach.

Testing the Implementation

Robust testing is essential to verify correctness across edge cases. The following test suite covers empty states, odd and even counts, duplicate values, and negative numbers:

import unittest


class TestMedianFinder(unittest.TestCase):
    def test_single_element(self):
        finder = MedianFinder()
        finder.add_num(1)
        self.assertEqual(finder.find_median(), 1.0)

    def test_two_elements(self):
        finder = MedianFinder()
        finder.add_num(1)
        finder.add_num(2)
        self.assertEqual(finder.find_median(), 1.5)

    def test_odd_count(self):
        finder = MedianFinder()
        for num in [5, 15, 1, 3]:
            finder.add_num(num)
        self.assertEqual(finder.find_median(), 4.0)

    def test_even_count(self):
        finder = MedianFinder()
        for num in [5, 15, 1, 3, 8]:
            finder.add_num(num)
        self.assertEqual(finder.find_median(), 5.0)

    def test_duplicates(self):
        finder = MedianFinder()
        for num in [2, 2, 2, 2]:
            finder.add_num(num)
        self.assertEqual(finder.find_median(), 2.0)

    def test_negative_numbers(self):
        finder = MedianFinder()
        for num in [-5, -10, -3, -7]:
            finder.add_num(num)
        self.assertEqual(finder.find_median(), -6.0)

    def test_mixed_positive_negative(self):
        finder = MedianFinder()
        for num in [-1, 2, -3, 4]:
            finder.add_num(num)
        self.assertEqual(finder.find_median(), 0.5)


if __name__ == "__main__":
    unittest.main()

Best Practices and Common Pitfalls

Handle Empty State Gracefully

Calling find_median before any numbers have been added will raise an IndexError when accessing self.lower_half[0]. In production code, you should either raise a descriptive custom exception or return a sentinel value:

def find_median(self):
    if not self.lower_half:
        raise ValueError("No numbers have been added yet")
    if len(self.lower_half) > len(self.upper_half):
        return float(-self.lower_half[0])
    return (-self.lower_half[0] + self.upper_half[0]) / 2.0

Remember the Negation Trick

A common mistake when implementing the max-heap in Python is forgetting to negate values during insertion and again when reading them back. Every value pushed to lower_half must be negated, and every value read or moved from lower_half must be negated again to restore the original number. Forgetting this leads to silently incorrect medians that can be difficult to debug.

Choose the Right Heap for the Extra Element

The implementation above stores the extra element (when the total count is odd) in the max-heap. This is a convention, not a requirement. You could equally store it in the min-heap, but you must be consistent throughout the code. Mixing conventions leads to off-by-one errors in median calculation.

Consider Integer Overflow in Other Languages

While Python handles arbitrarily large integers natively, the negation trick can cause overflow in languages with fixed-width integers. If porting this solution to Java, C++, or Go, use a max-heap implementation directly or ensure the numeric type can accommodate the negated values.

Use Type Hints for Clarity

In a professional codebase, adding type hints improves readability and enables static analysis tools to catch errors early:

from typing import List


class MedianFinder:
    def __init__(self) -> None:
        self.lower_half: List[int] = []
        self.upper_half: List[int] = []

    def add_num(self, num: int) -> None:
        if not self.lower_half or num <= -self.lower_half[0]:
            heapq.heappush(self.lower_half, -num)
        else:
            heapq.heappush(self.upper_half, num)

        if len(self.lower_half) > len(self.upper_half) + 1:
            heapq.heappush(self.upper_half, -heapq.heappop(self.lower_half))
        elif len(self.upper_half) > len(self.lower_half):
            heapq.heappush(self.lower_half, -heapq.heappop(self.upper_half))

    def find_median(self) -> float:
        if not self.lower_half:
            raise ValueError("No numbers have been added yet")
        if len(self.lower_half) > len(self.upper_half):
            return float(-self.lower_half[0])
        return (-self.lower_half[0] + self.upper_half[0]) / 2.0

Extending the Solution

Supporting a Sliding Window Median

A common variant asks for the median within a sliding window of fixed size. This requires the ability to remove arbitrary elements from the heaps, which standard heaps do not support efficiently. The typical approach uses lazy deletion: mark elements for removal and skip them when they surface at the top of a heap. This keeps amortized complexity manageable while adding implementation complexity.

Handling Floating-Point Streams

The same two-heap structure works for floating-point numbers without modification. However, be mindful of floating-point precision when comparing values. For most applications, direct comparison is sufficient, but in sensitive numerical contexts, consider using a tolerance-based comparison.

Parallel and Distributed Streams

For extremely high-volume streams, a single process may not keep up. In distributed settings, approximate median algorithms such as t-digest or the count-min sketch provide near-accurate medians with bounded memory. These are beyond the scope of this tutorial but are worth exploring when exact computation becomes infeasible.

Conclusion

The Find Median from Data Stream problem is a beautiful demonstration of how choosing the right data structure transforms an apparently expensive operation into an efficient one. By maintaining two heaps—a max-heap for the lower half and a min-heap for the upper half—we achieve O(log n) insertions and O(1) median queries, a substantial improvement over the naive sorting approach. Python's built-in heapq module, combined with the negation trick for max-heap behavior, makes the implementation concise and readable. Whether you are preparing for a coding interview or building a real-time analytics pipeline, mastering this pattern equips you with a powerful tool for streaming data processing. Remember to handle edge cases, test thoroughly, and consider extensions like sliding windows or approximate algorithms as your use cases grow in complexity.