LeetCode #295: Find Median from Data Stream

Description:

Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.

For example,

[2,3,4], the median is 3

[2,3], the median is (2 + 3) / 2 = 2.5

Design a data structure that supports the following two operations:

• void addNum(int num) - Add a integer number from the data stream to the data structure.
• double findMedian() - Return the median of all elements so far.

Example:

addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3) 
findMedian() -> 2

Follow up:

1. If all integer numbers from the stream are between 0 and 100, how would you optimize it?
2. If 99% of all integer numbers from the stream are between 0 and 100, how would you optimize it?

Solution:

• Build two heaps: one max-heap which stores the smaller part of the data, while the other min-heap stores the larger part of the data
• Everytime when you add new data, compare the new data with the roots of the two heaps first and then decide which heap you want to update (update time complexity: log(N/2))
• The median value can be easily found based on the roots the two heaps
• We can call this method "Hourglass" algorithm

class MedianFinder:

    def __init__(self):
        """
        initialize your data structure here.
        """
        self.left = []
        self.right = []

    def heapify_bottomup(self, array, i, flag='min'):
        if i > 0:
            iparent = int((i-1)//2)
            im = i
            if flag == 'min':
                if array[iparent] > array[i]: im = iparent
            if flag == 'max':
                if array[iparent] < array[i]: im = iparent
            if i != im:
                array[im], array[i] = array[i], array[im]
                self.heapify_bottomup(array, iparent, flag=flag)


    def heapify(self, array, i, flag='min'):
        if i < len(array)-1:
            ileft = 2*i + 1
            iright = 2*i + 2
            im = i
            if ileft < len(array):
                if flag == 'min':        
                    if array[ileft] < array[im]:
                        im = ileft
                if flag == 'max':        
                    if array[ileft] > array[im]:
                        im = ileft
            if iright < len(array):
                if flag == 'min':        
                    if array[iright] < array[im]:
                        im = iright
                if flag == 'max':        
                    if array[iright] > array[im]:
                        im = iright
            if im != i:
                array[i], array[im] = array[im], array[i]            
                self.heapify(array, im, flag=flag)


    def heap_delete(self, array, flag='min'):
        array[0], array[-1] = array[-1], array[0]
        array.pop()
        self.heapify(array, 0, flag=flag)


    def heap_insert(self, array, new, flag='min'):
        array.append(new)
        self.heapify_bottomup(array, len(array)-1, flag=flag)


    def addNum(self, num: int) -> None:
        if len(self.left) == len(self.right):
            if len(self.left) > 0:
                if num <= self.right[0]:
                    self.heap_insert(self.left, num, flag='max')
                else:
                    self.heap_insert(self.right, num, flag='min')
            else:
                self.left.append(num)

        elif len(self.left) > len(self.right):
            if num >= self.left[0]:
                self.heap_insert(self.right, num, flag='min')
            else:
                self.heap_insert(self.left, num, flag='max')
                self.heap_insert(self.right, self.left[0], flag='min')
                self.heap_delete(self.left, flag='max')
        else:
            if num <= self.right[0]:
                self.heap_insert(self.left, num, flag='max')
            else:
                self.heap_insert(self.right, num, flag='min')
                self.heap_insert(self.left, self.right[0], flag='max')
                self.heap_delete(self.right, flag='min')

        #print(self.left, self.right)


    def findMedian(self) -> float:
        if len(self.left) == len(self.right):
            if len(self.left) > 0:
                return (self.left[0] + self.right[0])/2
            else:
                return None
        elif len(self.left) > len(self.right):
            return self.left[0]
        else:
            return self.right[0]


# Your MedianFinder object will be instantiated and called as such:
# obj = MedianFinder()
# obj.addNum(num)
# param_2 = obj.findMedian()