Median from Data Stream — Amazon Interview Question

Solution using heaps

Photo by Nick Hillier on Unsplash

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

From a stream of numbers return the current median.

Understanding Median

To find the median we sort the elements and find the middle one (or the average of two middle elements). The below diagram is an illustration with an example of how we find the median.

Since we will be sorting the current elements this will give a time complexity of n² (with in-order insertion).

Insight from the median definition

One thing to notice here in the above diagram is that we are only interested in the middle element(s). And if we see it more closely we see that we are interested in the max element of the left side array and the min element from the right side of the array.

Heap is a data structure that maintains max or min at the top in O(logN). So we create two heaps, one for min-heap and another one for max-heap. The only constraint that we impose is that those two heaps, at max, only differ in size by 1. To balance the size of the heaps we pop the max/min element from the heap having a larger size and add that element to the min/max-heap of another heap.

To explain the balancing act of the heap below illustration is helpful.

Continuing with the remaining elements as they arrive.

Code Implementation

import heapq
class MedianFinder:
    def __init__(self):
        """
        initialize your data structure here.
        """
        self.left_heap = []
        self.right_heap = []
        self.size = 0
    def addNum(self, num: int) -> None:
        self.size += 1
        if self.size == 1:
            heapq.heappush(self.right_heap, num)
            return
        right_min = self.right_heap[0]
        
        if len(self.right_heap) > len(self.left_heap):
            if num < right_min:
                heapq.heappush(self.left_heap, -num)
            else:
                heapq.heappush(
                    self.left_heap, 
                    -heapq.heappushpop(self.right_heap, num)
                )
        else:
            if num > right_min:
                heapq.heappush(self.right_heap, num)
            else:
                heapq.heappush(
                    self.right_heap, 
                    -heapq.heappushpop(self.left_heap, -num)
                )
        
    def findMedian(self) -> float:
        if self.size % 2 == 0:
            return (-self.left_heap[0]+self.right_heap[0])/2
        else:
            return self.right_heap[0]

Hope it was helpful!

More content at plainenglish.io. Sign up for our free weekly newsletter. Get exclusive access to writing opportunities and advice in our community Discord.