Python O(lgn) using two heapq data sturctures


  • 3
    Z
    class MedianFinder:
    import heapq
    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.small = [] # store the small half, top is the largest in the small part
        self.large = [] # store the large half, top is the smallest in the large part
    
    def addNum(self, num):
        """
        Adds a num into the data structure.
        :type num: int
        :rtype: void
        """
        if len(self.small) == 0:
            heapq.heappush(self.small, -num)
            return
        if num <= -self.small[0]:
            # push to small part
            heapq.heappush(self.small, -num)
        else:
            # push to large part
            heapq.heappush(self.large, num)
        # adjust small and large balance
        if len(self.small) - len(self.large) == 2:
            heapq.heappush(self.large, -heapq.heappop(self.small))
        elif len(self.small) - len(self.large) == -2:
            heapq.heappush(self.small, -heapq.heappop(self.large))
    
    def findMedian(self):
        """
        Returns the median of current data stream
        :rtype: float
        """
        if len(self.small) == len(self.large):
            return (self.large[0] - self.small[0])/2.0
        return -float(self.small[0]) if len(self.small) > len(self.large) else float(self.large[0])

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.