Find Median in a Stream

In this class, We discuss Find median in a Stream.

Readers can prepare an entire competitive coding course to crack product development companies. Click Here.

The reader should have basic coding skills to work with competitive coding. Click here.

Question:

Given an input stream.

Find the median for each element in the input stream.

Time complexity: O(NlogN)

Space Complexity: O(N)

Example:

Input: [5, 15, 1, 3]

Output: [5, 10, 5, 4}

The median is the middle element in the list of sorted elements.

If there is an even number of numbers, then the median is the average of the middle two numbers.

Given three functions.

Insert heap function

Balance heap function

Find median function

The question says to use a heap data structure.

We can do the example using insertion sort. but time complexity is O(N^2)

Logic:

The step-by-step explanation is provided in the video.

Understanding this example helps a lot in using the heap data structure.

Code:

import heapq

class Solution:

    def __init__(self):

        self.small,self.large = [],[] #–> Max Heap = small , Min Heap = large

    def balanceHeaps(self):

        #Balance the two heaps size , such that difference is not more than one.

        if self.small and self.large and (-self.small[0]>self.large[0]):

            val = -heapq.heappop(self.small)

            heapq.heappush(self.large,val)

        if len(self.small)-len(self.large)>1:

            val = -heapq.heappop(self.small)

            heapq.heappush(self.large,val)

        elif len(self.large)-len(self.small)>1:

            val = -heapq.heappop(self.large)

            heapq.heappush(self.small,val)

    def getMedian(self):

        # return the median of the data received till now.

        if len(self.small)>len(self.large):

            return -self.small[0]

        elif len(self.small)<len(self.large):

            return self.large[0]

        else:

            return (-1*self.small[0]+self.large[0])/2

    def insertHeaps(self,x):

        #:param x: value to be inserted

        #:return: None

        heapq.heappush(self.small,-x)

        self.balanceHeaps()

a=[5,15,1,3]

ob=Solution()

k=[]

for i in a:

    ob.insertHeaps(i)

    z=ob.getMedian()

    k.append(z)

print(k)