295.Find Median from Data Stream

Tags: [data_structure], [stream], [streaming], [heap]

Com: {g}

Link: https://leetcode.com/problems/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.

Examples:

[2,3,4], the median is3

[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.

For example:

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

Solution: Heap

import heapq

class MedianFinder(object):

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


    def addNum(self, num):
        """
        :type num: int
        :rtype: void
        """
        if not self.max_heap or num <= self.max_heap[0][1]:
            heapq.heappush(self.max_heap, (-num, num))
        else:
            heapq.heappush(self.min_heap, (num, num))

        # balance the two heaps.
        while abs(len(self.max_heap) - len(self.min_heap)) > 1:
            if len(self.max_heap) > len(self.min_heap):
                _, elem = heapq.heappop(self.max_heap)
                heapq.heappush(self.min_heap, (elem, elem))
            else:
                _, elem = heapq.heappop(self.min_heap)
                heapq.heappush(self.max_heap, (-elem, elem))


    def findMedian(self):
        """
        :rtype: float
        """
        if not self.max_heap and not self.min_heap:
            return 0.0

        if (len(self.max_heap) + len(self.min_heap)) % 2 == 0:
            return (self.max_heap[0][1] + self.min_heap[0][1]) / 2.0
        else:
            return float(self.max_heap[0][1] if len(self.max_heap) > len(self.min_heap) else self.min_heap[0][1])


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

Note:

• Time complexity of init = O(1).
• Time complexity of addNum = O(lg(n/2)), n is the total number of elements at that moment. Because each time, once we find the abs(len(self.maxheap) - len(self.minheap)) > 1 we do balance, so at most the difference of size between two heap is 2, which means at the worst case, we just need do one heap pop and heap push then we can balance two heaps.
• Time complexity of findMedian = O(lg(n/2)), n is the total number of elements at that moment.