Median from Data Stream - Streaming Statistics Problem

The median from data stream problem requires maintaining a running median as new elements arrive, which can be efficiently solved using two heaps to partition the data around the median value.

1. Pattern Restatement

Clarifying questions to ask the interviewer:

1. What is the expected size of the data stream? Will it be bounded or unbounded?
2. Are duplicates allowed in the input data stream?
3. Should we return the median as a double or can we return an integer when the median is a whole number?
4. Can we modify the input data structure, or must we maintain immutability?
5. What is the target time complexity for the median retrieval operation? Should we aim for O(1) or is O(log n) acceptable?

Restate the pattern or problem type in your own words: The median from data stream is a classic streaming statistics problem where we need to efficiently maintain and retrieve the median value from a continuous stream of numbers. The median is the middle value in a sorted list, or the average of the two middle values if the list has an even number of elements.

Clarify the signal words and constraint shape that point to it: Key signals include "data stream," "real-time updates," "running median," "dynamic median," and "online algorithm." The constraint shape involves handling an unbounded sequence of numbers where we must support both insertion and median retrieval operations efficiently.

Note any assumptions you are making: I'm assuming that the data stream is unbounded and we're processing numbers one at a time. I also assume that we want to support both adding new elements and retrieving the current median efficiently.

Identify whether the pattern is best suited for static, streaming, or online scenarios if relevant: This pattern is specifically designed for streaming/online scenarios where data arrives continuously and we must maintain statistics in real-time. It's not suitable for static datasets where all data is known upfront, as there are more efficient approaches for batch processing.

The core insight is that we need to maintain a data structure that allows us to quickly find the median without having to re-sort or recompute everything every time a new element arrives.

2. When to Use This Pattern

Signal words and constraint shapes:

• "median" or "middle element" in a dynamic dataset
• "data stream" indicating continuous input
• "real-time" or "online" processing requirements
• Need to maintain median without sorting the entire dataset each time
• Input size constraints that make full sorting impractical (e.g., millions of elements)

Why this pattern fits better than alternatives: The median-finding pattern is specifically designed for maintaining a running median in a data stream. Compared to alternatives like:

• Simple sorting approach: O(n log n) per query, too slow for large datasets
• Hash-based approaches: Don't provide ordering information needed for median calculation
• Single heap approaches: Can't efficiently maintain both halves of the data

The key insight is that we need to partition data into two halves - one containing elements less than or equal to the median, and another containing elements greater than or equal to the median. This naturally maps to using two heaps (max-heap and min-heap) to maintain the partitioning property.

When variations apply:

• For kth smallest/largest elements, use a single heap with size tracking
• For multiple medians (e.g., sliding window), consider using a combination of heaps with additional tracking structures
• For static datasets, simple sorting or quickselect is more appropriate

Core insight: The pattern works by maintaining two heaps that act as a "balanced partition" - the max-heap stores the smaller half and the min-heap stores the larger half. The top elements of these heaps give us direct access to the median without needing to sort or reorganize the entire dataset.

3. Approach Evolution

Interview communication script:

My first instinct is to use a brute force approach where I maintain the entire array and sort it every time we need the median. But that's O(n² log n) because for each insertion, I have to sort the entire array. Let me think about how to optimize this.

I notice that I don't actually need the entire array sorted — I just need the middle element(s). If I could somehow keep track of the elements in a way that allows me to quickly access the median, that would be better.

I think I can use two heaps to maintain the elements in a way that the median is always accessible at the top of one of the heaps. If I maintain a max heap for the smaller half and a min heap for the larger half, then:

• The top of max heap gives me the largest element in smaller half
• The top of min heap gives me the smallest element in larger half

This way, I can get the median in O(1) time and maintain it with O(log n) insertions.

Let me think about the invariants I need to maintain:

1. The max heap contains at most one more element than the min heap
2. All elements in max heap are ≤ all elements in min heap

This approach reduces the complexity from O(n² log n) to O(n log n) time and O(n) space.

Brute-force approach and why it is not ideal: The brute force approach would involve maintaining a sorted array and inserting new elements in the correct position to keep it sorted. For each insertion, we'd need to:

1. Find the correct position using binary search - O(log n)
2. Shift elements to make space - O(n)
3. Insert the new element - O(1)

This gives us O(n²) time complexity for n insertions, which is inefficient.

Improved approach and the key optimization: The key insight is to use two heaps instead of maintaining a sorted array. By using a max heap for the smaller half and min heap for the larger half, we can achieve O(log n) insertion time while maintaining O(1) median access.

Final approach with step-by-step reasoning:

1. Use two heaps: max heap (smaller half) and min heap (larger half)
2. Keep the heaps balanced so max heap size is either equal to or one more than min heap size
3. When inserting a new element:
   • Add it to the appropriate heap based on comparison with max heap's top
   • Rebalance heaps if needed to maintain size constraint
4. The median is accessible in O(1) time from the top of one or both heaps

Important invariants that must hold:

1. Size constraint: max heap size ≤ min heap size + 1
2. Ordering constraint: all elements in max heap ≤ all elements in min heap
3. Balance property: the difference in sizes between heaps is at most 1

4. Data Structure Choice

For the median from data stream problem, I need to maintain a dynamic collection of numbers that supports efficient insertion and median retrieval operations.

Primary Data Structures:

1. Two Heaps (Max-heap and Min-heap) - This is the optimal approach for this problem. I'll use:

   • A max-heap to store the smaller half of numbers (root contains the largest among smaller half)
   • A min-heap to store the larger half of numbers (root contains the smallest among larger half)

2. Heap Properties:

   • Max-heap (smaller half): root = max of all elements in this heap
   • Min-heap (larger half): root = min of all elements in this heap

Why This Choice Works:

• The heaps naturally maintain the partition property where all elements in max-heap ≤ all elements in min-heap
• Insertion is O(log n) time complexity using heap operations
• Median retrieval is O(1) - just peek at the root(s)
• Space complexity is O(n) to store all elements

Time and Space Impact:

• Time Complexity: O(log n) for insertion, O(1) for median retrieval
• Space Complexity: O(n) to store all elements across both heaps

Helper Structures:

• Java's PriorityQueue with custom comparators to simulate max-heap and min-heap behavior
• No additional helper structures needed beyond the two heaps

This approach is superior to maintaining a sorted array or list because:

• Sorted arrays require O(n) insertion due to shifting elements
• Sorted lists have similar issues with insertion cost
• Heaps provide the optimal balance of insertion and retrieval efficiency

The key insight is that we don't need to maintain a complete sorted order - we only need to know the middle element(s), which can be efficiently tracked using the two-heap partitioning strategy.

5. Algorithm Walkthrough

Let me trace through the algorithm with a concrete example to demonstrate how it works.

Example: Stream of numbers: [5, 15, 1, 3]

Step 1: Add 5

• Left heap (max): [5]
• Right heap (min): []
• Median: 5.0

Step 2: Add 15

• Left heap (max): [5]
• Right heap (min): [15]
• Balanced heaps: left size = 1, right size = 1
• Median: (5 + 15) / 2 = 10.0

Step 3: Add 1

• Left heap (max): [5, 1] → [5]
• Right heap (min): [15, 1] → [15]
• Left size = 1, right size = 1
• Median: (5 + 15) / 2 = 10.0

Step 4: Add 3

• Left heap (max): [5, 1] → [5, 3] → [5]
• Right heap (min): [15, 1] → [15, 3] → [15]
• Left size = 2, right size = 2
• Median: (5 + 15) / 2 = 10.0

The key insight is that we maintain the invariant where the left heap contains the smaller half of elements (with max at root) and right heap contains larger half (with min at root). When we add a new element, we decide which heap to place it in based on its value relative to the current medians, then rebalance if needed.

The algorithm correctly handles all edge cases including odd and even-length streams, and maintains the mathematical definition of median throughout the process.

6. Java Solution

import java.util.PriorityQueue;

public class MedianFinder {
    private PriorityQueue<Integer> maxHeap; // stores smaller half
    private PriorityQueue<Integer> minHeap; // stores larger half
    
    public MedianFinder() {
        // Max heap: custom comparator to make it a max heap
        maxHeap = new PriorityQueue<>((a, b) -> b - a);
        // Min heap: default behavior
        minHeap = new PriorityQueue<>();
    }
    
    public void addNum(int num) {
        // Add to max heap first (smaller half)
        maxHeap.offer(num);
        
        // Balance: move largest element from max heap to min heap
        if (!maxHeap.isEmpty() && !minHeap.isEmpty() && 
            maxHeap.peek() > minHeap.peek()) {
            int max = maxHeap.poll();
            minHeap.offer(max);
        }
        
        // Rebalance heaps to maintain size difference at most 1
        if (maxHeap.size() > minHeap.size() + 1) {
            minHeap.offer(maxHeap.poll());
        } else if (minHeap.size() > maxHeap.size() + 1) {
            maxHeap.offer(minHeap.poll());
        }
    }
    
    public double findMedian() {
        if (maxHeap.size() == minHeap.size()) {
            // Even number of elements: average of both heaps' tops
            return (maxHeap.peek() + minHeap.peek()) / 2.0;
        } else if (maxHeap.size() > minHeap.size()) {
            // Odd number of elements: max heap has the extra element
            return (double) maxHeap.peek();
        } else {
            // Odd number of elements: min heap has the extra element
            return (double) minHeap.peek();
        }
    }
}

This solution uses two heaps to maintain the data stream in a way that allows O(1) median retrieval. The max heap holds the smaller half of numbers, and the min heap holds the larger half. By maintaining size balance between the heaps, we ensure that the median can be found in constant time by examining the top elements of one or both heaps.

7. Complexity Analysis

Time Complexity: O(log n) per insertion operation

Each addNum operation requires:

• Inserting into a heap: O(log n)
• Potentially rebalancing heaps: O(1)
• Finding median: O(1)

The findMedian operation is O(1) as it only accesses the top elements of the heaps.

Space Complexity: O(n)

We store all elements in two heaps, so space grows linearly with the number of elements added. Each element is stored exactly once across both heaps.

Why these bounds are correct:

The heap operations (insertion and deletion) have logarithmic time complexity, which dominates the overall cost. Since we maintain at most one element difference between heaps, the total number of elements in both heaps is always 2n where n is the count of elements added. The median access is constant time because we only need to peek at the top elements.

Optimality:

This solution is near-optimal for the streaming median problem. The O(log n) insertion time matches the theoretical lower bound for maintaining a sorted structure with insertions, as any comparison-based approach must spend at least logarithmic time to maintain order. The space complexity is optimal since we must store all elements to compute the median.

The solution achieves the best possible time complexity for a streaming scenario where we cannot pre-sort the data, making it ideal for real-time systems where elements arrive one at a time.

8. Edge Cases & Pitfalls

Empty and Minimal Inputs

• Empty stream: The solution must handle zero elements gracefully, typically returning a default value or throwing an appropriate exception
• Single element: Should correctly return that element as the median
• Two elements: Must properly handle the average of two numbers

Duplicate Values and Repeated States

• Duplicate values: The solution should treat duplicates as valid elements in the median calculation - they don't change the algorithm's logic but may affect the final result
• Repeated insertion: The same value can be inserted multiple times, and the data structures must maintain correctness through all such operations

Boundary Conditions

• Integer overflow: When calculating averages of large integers, casting to long is necessary to prevent overflow
• Empty heap states: When one or both heaps become empty, the solution must handle these edge cases properly
• Heap size imbalance: The algorithm requires maintaining heap size differences of at most 1, so careful handling is needed when elements are added or removed

Common Implementation Mistakes

• Incorrect heap size tracking: Failing to maintain the invariant that heap sizes differ by at most 1
• Wrong median calculation: Not properly handling the case where we have an even number of elements (taking average vs. taking the lower element)
• Heap ordering confusion: Using PriorityQueue with default ordering instead of custom comparators to maintain max-heap and min-heap properties
• Not handling concurrent access: In multi-threaded environments, the solution must be thread-safe or properly synchronized
• Missing proper initialization: Not initializing heaps correctly when the first element is added

The solution must be robust enough to handle all these edge cases without compromising correctness or performance.

9. Alternative Approaches

There are several alternative approaches to solving the median from data stream problem, each with different trade-offs:

1. Sorting-based approach with binary search insertion: Instead of maintaining sorted order explicitly, we can use a TreeSet or manually maintain a sorted list using binary search to find insertion points. Each insertion takes O(log n) time for finding the position, but shifting elements to maintain sorted order still requires O(n) time in worst case. This approach has O(n log n) overall complexity for n insertions.

2. Two heaps (max-heap + min-heap): This is the most commonly discussed alternative and is actually the optimal approach for this problem. We maintain a max-heap for the smaller half of elements and a min-heap for the larger half. The heaps are kept balanced so that their sizes differ by at most one. Insertion takes O(log n) time and median retrieval is O(1). This approach is more complex to implement correctly but matches our chosen solution's time complexity.

3. Balanced binary search tree (e.g., TreeSet in Java): Using a self-balancing BST like TreeSet allows for O(log n) insertion and O(1) median access when we maintain the size of each subtree. However, Java's TreeSet doesn't provide direct access to the k-th element, so we'd need to implement our own tree structure or use TreeMap with additional tracking.

4. Simple array with linear search and sort: For very small data streams or when simplicity is preferred over performance, we could maintain an array and sort it after each insertion. This approach has O(n log n) complexity per insertion, making it unsuitable for large streams.

The two-heap approach is generally preferred over our solution because:

• It's more intuitive once understood
• It has clearer separation of concerns (heap management vs. median calculation)
• It's a well-known pattern that appears in many similar problems
• The implementation is more robust to edge cases

Our chosen solution with PriorityQueue and TreeMap is more efficient in terms of memory usage for the specific case where we only need to track counts, but the two-heap approach is more commonly taught and used in interviews.

10. Test Cases

import java.util.*;

public class MedianFinderTestCases {
    public static void main(String[] args) {
        // Test case 1: Basic functionality with odd number of elements
        MedianFinder mf1 = new MedianFinder();
        int[] nums1 = {5, 15, 1, 3};
        double[] expected1 = {5.0, 10.0, 5.0, 4.0};
        testMedianFinder(mf1, nums1, expected1);
        
        // Test case 2: Even number of elements
        MedianFinder mf2 = new MedianFinder();
        int[] nums2 = {1, 2, 3, 4};
        double[] expected2 = {1.0, 1.5, 2.0, 2.5};
        testMedianFinder(mf2, nums2, expected2);
        
        // Test case 3: Single element
        MedianFinder mf3 = new MedianFinder();
        int[] nums3 = {42};
        double[] expected3 = {42.0};
        testMedianFinder(mf3, nums3, expected3);
        
        // Test case 4: Empty input (should handle gracefully)
        MedianFinder mf4 = new MedianFinder();
        int[] nums4 = {};
        // No operations expected - just verify no exceptions
        
        // Test case 5: Negative numbers
        MedianFinder mf5 = new MedianFinder();
        int[] nums5 = {-1, -2, -3, -4};
        double[] expected5 = {-1.0, -1.5, -2.0, -2.5};
        testMedianFinder(mf5, nums5, expected5);
        
        // Test case 6: Duplicate elements
        MedianFinder mf6 = new MedianFinder();
        int[] nums6 = {1, 1, 2, 2};
        double[] expected6 = {1.0, 1.0, 1.5, 1.5};
        testMedianFinder(mf6, nums6, expected6);
        
        // Test case 7: Large numbers
        MedianFinder mf7 = new MedianFinder();
        int[] nums7 = {1000000, 2000000, 3000000};
        double[] expected7 = {1000000.0, 1500000.0, 2000000.0};
        testMedianFinder(mf7, nums7, expected7);
    }
    
    private static void testMedianFinder(MedianFinder mf, int[] nums, double[] expected) {
        System.out.println("Testing MedianFinder with input: " + Arrays.toString(nums));
        List<Double> results = new ArrayList<>();
        
        for (int num : nums) {
            mf.addNum(num);
            results.add(mf.findMedian());
        }
        
        System.out.println("Results: " + results);
        System.out.println("Expected: " - Arrays.toString(expected));
        
        // Verify results match expected values
        for (int i = 0; i < expected.length; i++) {
            if (Math.abs(results.get(i) - expected[i]) > 1e-9) {
                System.out.println("FAIL: Index " + i + " - got " + results.get(i) + ", expected " + expected[i]);
            }
        }
        System.out.println("Test completed successfully\n");
    }
}

Sample Test Cases:

1. Basic functionality: Input {5, 15, 1, 3} → Expected medians {5.0, 10.0, 5.0, 4.0}
2. Even count: Input {1, 2, 3, 4} → Expected medians {1.0, 1.5, 2.0, 2.5}
3. Single element: Input {42} → Expected median {42.0}
4. Negative numbers: Input {-1, -2, -3, -4} → Expected medians {-1.0, -1.5, -2.0, -2.5}
5. Duplicate elements: Input {1, 1, 2, 2} → Expected medians {1.0, 1.0, 1.5, 1.5}
6. Large numbers: Input {1000000, 2000000, 3000000} → Expected medians {1000000.0, 1500000.0, 2000000.0}

Tricky Cases:

• Empty input: No elements added, should handle gracefully
• Mixed positive/negative: Ensure proper ordering in heaps
• Large input size: Verify O(log n) insertion performance

Boundary Cases:

• Single element: Edge case for median calculation
• Two elements: Simple average case
• Duplicate values: Heap balancing considerations
• Integer overflow: Large numbers within int range

11. Final Summary

The core idea is to maintain a running median from a data stream using two heaps to efficiently partition and access the middle elements. This solution works by keeping the smaller half in a max heap and the larger half in a min heap, ensuring O(1) median access with O(log n) insertion time. The key insight is that we only need to track the boundary elements (heaps' tops) to determine the median, making it highly efficient for streaming data scenarios where the entire dataset cannot be stored in memory.