Find Median from Data Stream | Brute Force to Optimal

1. Problem Statement

Imagine you are receiving a continuous stream of numbers (like sensor readings or financial ticks). At any point, someone might hit a "pause" button and ask you: "What is the median of all the numbers you've seen so far?"

What is a median? It's the middle value of a sorted list of numbers.
- If the list has an odd number of elements, the median is the exact middle number (e.g., in [1, 2, 3], median is 2).
- If the list has an even number of elements, the median is the average of the two middle numbers (e.g., in [1, 2, 3, 4], median is (2 + 3) / 2 = 2.5).

You need to design a system that supports two operations efficiently:

addNum(int num): Adds a new number from the stream into your data structure.
findMedian(): Returns the median of all numbers added so far as a double.

2. Brute Force Solution

Intuition

The most straightforward way to find the middle of a dataset is to keep all the numbers, sort them whenever asked, and then pick the middle element(s).

Algorithm

addNum: Simply append the new number to the end of an array.
findMedian: Sort the entire array. If the size is odd, return the middle element. If even, return the average of the two middle elements.

Dry Run

Stream: 5, 3

addNum(5): Array = [5]
addNum(3): Array = [5, 3]
findMedian(): Sort to [3, 5]. Size is even. Return (3 + 5) / 2.0 = 4.0.

C++ Implementation

#include <vector>
#include <algorithm>

class MedianFinderBrute {
    std::vector<int> nums;

public:
    void addNum(int num) {
        nums.push_back(num); // Just append
    }
    
    double findMedian() {
        std::vector<int> sorted_nums = nums; // Copy to avoid modifying original stream order
        std::sort(sorted_nums.begin(), sorted_nums.end());
        
        int n = sorted_nums.size();
        if (n % 2 == 1) {
            return sorted_nums[n / 2];
        } else {
            return (sorted_nums[n / 2 - 1] + sorted_nums[n / 2]) / 2.0;
        }
    }
};

3. Better Solution (Insertion Sort / Binary Search)

Intuition

Instead of appending blindly and sorting later, let's maintain a sorted array at all times. When a new number arrives, we figure out exactly where it belongs in our sorted array and insert it there.

Algorithm

addNum: Use Binary Search (via std::lower_bound in C++) to find the correct insertion index for the new number. Then, insert it.
findMedian: Since the array is always sorted, we just jump straight to the middle index (or indices) and calculate the median.

Dry Run

Stream: 5, 3, 8

addNum(5): Array = [5]
addNum(3): Binary search finds index 0. Insert. Array = [3, 5]
addNum(8): Binary search finds index 2. Insert. Array = [3, 5, 8]
findMedian(): Array is already sorted. Size 3 (odd). Return Array[1] = 5.0.

C++ Implementation

#include <bits/stdc++.h>
using namespace std;

class MedianFinder {
    vector<int> nums;   // Always stores elements in sorted order

public:

    // Add a new number to the data stream
    void addNum(int num) {

        // Find the position where 'num' should be inserted
        // so that the vector remains sorted.
        auto it = lower_bound(nums.begin(), nums.end(), num);

        // Insert the element at that position.
        // Elements after 'it' are shifted one position to the right.
        nums.insert(it, num);
    }

    // Return the median of all inserted numbers
    double findMedian() {

        int n = nums.size();

        // Odd number of elements
        if (n % 2 == 1) {
            return nums[n / 2];
        }

        // Even number of elements
        return (nums[n / 2 - 1] + nums[n / 2]) / 2.0;
    }
};

Complexity Analysis

Time Complexity: * addNum(): $O(n)$ . Binary search takes $O(\log n)$ , but std::vector::insert takes $O(n)$ because it shifts elements in memory.
- findMedian(): $O(1)$ . We just access indices.
Space Complexity: $O(n)$ to store the array.

Advantages & Disadvantages

Advantages: findMedian is now instantly $O(1)$ .
Disadvantages: addNum is now $O(n)$ . In a data stream with millions of elements, shifting arrays on every insertion is far too slow.

Why We Move to the Next Approach

We realized that keeping the entire array perfectly sorted is overkill. We only care about the middle element(s). What if we split the data in half, and just kept track of the "boundary" between the smaller half and the larger half?

4. Optimal Solution (Two Heaps)

Intuition

Imagine a sorted array divided perfectly down the middle: [1, 2, 3] | [4, 5, 6] Left Half | Right Half

To find the median, we only need to look at the largest number in the Left Half (which is 3) and the smallest number in the Right Half (which is 4). We don't care that 1 is smaller than 2, or 6 is bigger than 5. We only care about the elements hugging the center boundary.

What data structures are perfect for instantly fetching the largest or smallest element of a dataset? Heaps (Priority Queues).

Max Heap: Stores the smaller half of numbers. It gives us instant access to the largest number in this lower half.
Min Heap: Stores the larger half of numbers. It gives us instant access to the smallest number in this upper half.

Key Observations

To make this work, we must strictly enforce two rules (Invariants):

Ordering Invariant: Every number in the Max Heap (smaller half) must be less than or equal to every number in the Min Heap (larger half).
Heap Size Invariant (Balancing): The heaps must be balanced in size.
- If total elements are even, both heaps have the same size.
- If total elements are odd, we will arbitrarily let the Max Heap hold the extra element. Thus, MaxHeap.size() == MinHeap.size() + 1.

Algorithm

For addNum(num):

Add and sort: Always push the new number onto the maxHeap first.
Enforce Ordering: Because we blindly pushed to maxHeap, that number might actually belong in the upper half! To fix this, pop the top of the maxHeap and push it into the minHeap.
Enforce Size (Balance): The maxHeap is supposed to hold the extra element (if odd). If the minHeap suddenly has more elements than the maxHeap, pop the top of the minHeap and push it back to the maxHeap.

For findMedian():

If maxHeap.size() > minHeap.size(), the total count is odd. Return maxHeap.top().
Otherwise, sizes are equal (even count). Return (maxHeap.top() + minHeap.top()) / 2.0.

10. C++ Implementation

#include <bits/stdc++.h>
using namespace std;

class MedianFinder {
private:

    // Max Heap stores the smaller half of the numbers.
    // The largest element of the smaller half stays at the top.
    priority_queue<int> maxHeap;

    // Min Heap stores the larger half of the numbers.
    // The smallest element of the larger half stays at the top.
    priority_queue<int, vector<int>, greater<int>> minHeap;

public:

    MedianFinder() {}

    // Insert a new number into the data stream
    void addNum(int num) {

        // Step 1:
        // Insert the new number into the max heap.
        maxHeap.push(num);

        // Step 2:
        // Move the largest element from the max heap
        // to the min heap to maintain the ordering property.
        minHeap.push(maxHeap.top());
        maxHeap.pop();

        // Step 3:
        // Ensure maxHeap has either the same number of
        // elements as minHeap or exactly one more.
        if (minHeap.size() > maxHeap.size()) {
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }
    }

    // Return the median of all inserted numbers
    double findMedian() {

        // If maxHeap has one extra element,
        // the median is simply its top.
        if (maxHeap.size() > minHeap.size()) {
            return maxHeap.top();
        }

        // Otherwise, the median is the average
        // of the two middle elements.
        return (maxHeap.top() + minHeap.top()) / 2.0;
    }
};

[Max Heap]             [Min Heap]
      (Lower Half)          (Upper Half)
          5                       7
         / \                     / \
        3   2                   8   9
       /
      1

MaxHeap.top() = 5
MinHeap.top() = 7
Total size = 7 (Odd). Median is MaxHeap.top() = 5.0

+------+-----+---------------------------------------------------------------+----------------------+----------------------+-------------------+
|Step  | Num | addNum() Actions (Push → Transfer → Balance)                  | Max Heap (Lower)     | Min Heap (Upper)     | findMedian()      |
+------+-----+---------------------------------------------------------------+----------------------+----------------------+-------------------+
| 1    |  5  | Push to Max → Transfer Max→Min → Move Min→Max                 | [5]                  | []                   | 5.0               |
| 2    |  3  | Push to Max → Transfer Max top(5)→Min → Balanced              | [3]                  | [5]                  | (3+5)/2 = 4.0     |
| 3    |  8  | Push to Max → Transfer Max top(8)→Min → Move Min top(5)→Max   | [5,3]                | [8]                  | 5.0               |
| 4    |  9  | Push to Max → Transfer Max top(9)→Min → Balanced              | [5,3]                | [8,9]                | (5+8)/2 = 6.5     |
| 5    |  2  | Push to Max → Transfer Max top(5)→Min → Move Min top(5)→Max   | [5,3,2]              | [8,9]                | 5.0               |
| 6    |  1  | Push to Max → Transfer Max top(5)→Min → Balanced              | [3,2,1]              | [5,8,9]              | (3+5)/2 = 4.0     |
| 7    |  7  | Push to Max → Transfer Max top(7)→Min → Move Min top(5)→Max   | [5,3,2,1]            | [7,8,9]              | 5.0               |
+------+-----+---------------------------------------------------------------+----------------------+----------------------+-------------------+

6. Correctness Proof

Why does this always work?

The Median Boundary: The median of a sorted array relies entirely on the middle elements.
Invariants Maintained: Because every insertion pushes an element through the Max Heap and filters it into the Min Heap, the largest element of the lower half naturally bubbles up to the Max Heap's top, and the smallest element of the upper half sits at the Min Heap's top.
By strictly enforcing the rule that MaxHeap.size() is either equal to MinHeap.size() or larger by exactly 1, we guarantee that the "middle" elements are always at the tops of these two heaps.

7. Complexity Analysis

Time Complexity:
- addNum(): $O(\log n)$ . Pushing or popping from a heap of size $n/2$ takes logarithmic time. In the worst case, we do 3 heap operations (push Max, push Min/pop Max, push Max/pop Min). $3 \times O(\log (n/2))$ simplifies perfectly to $O(\log n)$ .
- findMedian(): $O(1)$ . Accessing the top of a heap is a constant-time operation.
Space Complexity: $O(n)$ . We store every element from the stream exactly once across the two heaps.

8. Edge Cases

Empty Stream: LeetCode bounds usually guarantee findMedian won't be called on an empty stream. If it were, you would throw an exception or return a sentinel value.
One Element: The first element naturally falls into the Max Heap. findMedian safely returns it without dividing by zero or accessing an empty Min Heap.
Duplicate Values: Heaps handle duplicates naturally without any special logic.
Negative Numbers: The sorting logic (less than / greater than) of heaps handles negative numbers flawlessly.
Extremely Large Inputs: Priority Queues scale well in memory up to system limits. For massive streams exceeding RAM, you would need an external approximation algorithm like a Count-Min Sketch or T-Digest, but that is out of scope for standard DSA interviews.

9. Common Interview Mistakes

Overcomplicating the balancing logic: Many candidates try to use if (num < maxHeap.top()) to decide which heap to insert into. This requires writing multiple complex, bug-prone if/else branches to handle edge cases. The "Push to Max, transfer to Min, balance back" logic avoids conditionals entirely.
Incorrect Average Calculation: Writing (maxHeap.top() + minHeap.top()) / 2 (integer division) instead of / 2.0 (floating point). In C++, integer division truncates decimals, causing failures on test cases like [1, 2] expecting 1.5.
Forgetting STL Syntax: In C++, std::priority_queue is a Max Heap by default. Remembering the verbose syntax for a Min Heap is critical: std::priority_queue<int, std::vector<int>, std::greater<int>>

video reference:

reference video

problem link:

https://leetcode.com/problems/find-median-from-data-stream/description/

1. Problem Statement

2. Brute Force Solution

3. Better Solution (Insertion Sort / Binary Search)

4. Optimal Solution (Two Heaps)

Algorithm

6. Correctness Proof

7. Complexity Analysis

8. Edge Cases

9. Common Interview Mistakes

chakradhar