Maintain the Median of a Data Stream

Question

This question asks the candidate to design a system that can efficiently maintain the median of a continuously incoming data stream.

Codemia · Accepted Answer

To maintain the median of a data stream, we can use two heaps: a max-heap for the lower half of the numbers and a min-heap for the upper half. This allows us to efficiently balance the two heaps to ensure that the median can be retrieved in constant time.

1. Initialize two heaps: a max-heap for the lower half and a min-heap for the upper half.
2. For each incoming number:
   - If the number is less than or equal to the max of the lower half, add it to the max-heap.
   - Otherwise, add it to the min-heap.
3. Balance the heaps if necessary:
   - If the max-heap has more than one extra element than the min-heap, move the max element from the max-heap to the min-heap.
   - If the min-heap has more elements than the max-heap, move the min element from the min-heap to the max-heap.
4. The median is the max of the max-heap if the heaps are of unequal size, or the average of the max of the max-heap and the min of the min-heap if they are of equal size.

The time complexity for inserting a new number is O(log n) due to the heap operations, where n is the number of elements in the heaps. Retrieving the median is O(1). Therefore, the overall complexity for processing a stream of k numbers is O(k log k).

```python
import heapq

class MedianFinder:
    def __init__(self):
        self.low = []  # max-heap (inverted min-heap)
        self.high = []  # min-heap

def addNum(self, num: int) -> None:
        heapq.heappush(self.low, -num)
        # Ensure max of low is less than min of high
        if (self.low and self.high and (-self.low[0] > self.high[0])):
            heapq.heappush(self.high, -heapq.heappop(self.low))

# Balance the heaps if necessary
        if len(self.low) > len(self.high) + 1:
            heapq.heappush(self.high, -heapq.heappop(self.low))
        elif len(self.high) > len(self.low):
            heapq.heappush(self.low, -heapq.heappop(self.high))

def findMedian(self) -> float:
        if len(self.low) > len(self.high):
            return float(-self.low[0])
        return (-self.low[0] + self.high[0]) / 2.0
```

Consider edge cases such as an empty stream or a stream with repeated numbers. The algorithm should handle these gracefully, ensuring that the median is calculated correctly even when the input is not varied.

Maintain the Median of a Data Stream

Xiaopeng

What the Interviewer Expects

Key Topics to Cover

How to Approach This

Possible Follow-up Questions

Sharpen Your Skills on Codemia

Sample Answer

Approach

Algorithm

Submit Your Answer

Xiaopeng Software Engineer Interview Guide

Related Questions

Binary Search on graph

Count kth largest element in array

Find connected components in string

Merge Three Sorted strings

Binary Search on string