Find Median From Data Stream - Leetcode Solution

Problem:

The problem "Find Median From Data Stream" asks us to design a class that takes a stream of integers as input and returns the median of all the integers in the stream at any given time. The class should be able to handle dynamic insertion of new integers in the stream. The median is defined as the middle element when the integers are sorted in ascending order. If there are an even number of integers in the stream, the median is the average of the two middle elements.

Solution:

The solution to this problem involves finding an efficient way to insert new integers into the stream and retrieve the median at any given time. One efficient approach is to use two heaps to track the median.

We can use a max-heap and a min-heap to keep track of the larger half and smaller half of the stream, respectively. The max-heap holds the smaller half of the stream, and the min-heap holds the larger half. We maintain the property that the size of the max-heap is either equal to or one greater than the size of the min-heap. This ensures that the median can be easily computed as the maximum value in the max-heap when their sizes are equal, and as the average of the maximum value in the max-heap and the minimum value in the min-heap when their sizes differ by one.

To insert a new integer into the stream, we first compare it with the maximum value in the max-heap. If the new integer is less than or equal to the maximum value, we insert it into the max-heap. Otherwise, we insert it into the min-heap. We then balance the size of the heaps so that their sizes differ by at most one.

To retrieve the median at any given time, we check the sizes of the max-heap and the min-heap. If they are equal, we return the maximum value in the max-heap. If their sizes differ by one, we return the average of the maximum value in the max-heap and the minimum value in the min-heap.

Here is the implementation of the class using two heaps in Python:

import heapq

class MedianFinder:
    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.max_heap = []   # contains smaller half of the stream
        self.min_heap = []   # contains larger half of the stream

    def addNum(self, num: int) -> None:
        if len(self.max_heap) == len(self.min_heap):
            if not self.min_heap or num <= self.min_heap[0]:
                heapq.heappush(self.max_heap, -num)
            else:
                heapq.heappush(self.max_heap, -heapq.heappushpop(self.min_heap, num))
        else:
            if num >= -self.max_heap[0]:
                heapq.heappush(self.min_heap, num)
            else:
                heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))
                heapq.heappush(self.max_heap, -num)

    def findMedian(self) -> float:
        if len(self.max_heap) == len(self.min_heap):
            return (-self.max_heap[0] + self.min_heap[0]) / 2
        else:
            return -self.max_heap[0]

In the addNum method, we first check if the size of the two heaps is equal. If so, we compare the new number with the minimum value in the min-heap. If the new number is less than or equal to the minimum value, we insert it into the max-heap. Otherwise, we insert it into the min-heap after first popping the minimum value from the min-heap and pushing it into the max-heap. If the size of the two heaps is not equal, we compare the new number with the maximum value in the max-heap. If the new number is greater than or equal to the maximum value, we insert it into the min-heap. Otherwise, we swap the maximum value of the max-heap with the new number and insert the swapped value into the min-heap. We then balance the size of the two heaps so that their sizes differ by at most one.

In the findMedian method, we first check if the size of the two heaps is equal. If so, we return the average of the maximum value in the max-heap and the minimum value in the min-heap. If the size of the two heaps is not equal, we return the maximum value in the max-heap.

1