Python Log(n) Based Solution Using 2 heaps


  • 0
    G

    My solution uses two heaps (a low and high heap), and then gets the median by either averaging the largest element in the small heap with the smallest in the large heap (if there is an even number of elements) or taking the largest element from the small heap or the smallest element from the large heap (if there are an odd number of elements).

    Inserting happens in log(n) time too, since we have to add the element to one of the heaps and potentially rebalance the two heaps in order to ensure the difference between the number of elements in the two heaps is less than 1.

    There is also a great explanation on coursera found here at about the 10:20 mark in the video:

    https://class.coursera.org/algo-003/lecture/62

    class MedianFinder:
        def __init__(self):
            """
            Initialize your data structure here.
            """
            self.lCount = 0
            self.hCount = 0
    
            self.h_low = [] #a max heap of the n/2 lowest numbers 
            self.h_high = [] #a min heap of the n/2 highest numbers
            
    
        def addNum(self, num):
            #print "adding {0} to the list".format(newElement)
            newElement = num
            loMax = hiMin = None
            if self.lCount >0:
                loMax = -1*self.h_low[0] #the max element of the lower half
            if self.hCount >0:
                hiMin = self.h_high[0] #the min element of the upper half
    
            if loMax ==None and hiMin == None: #both are empty
                self._add_to_low(newElement)
                return
            #they cannot both be empty
            if loMax:
                if newElement < loMax:
                    self._add_to_low(newElement)
                else:
                    self._add_to_high(newElement)
            else:
                if newElement >= hiMin:
                    self._add_to_high(newElement)
                else:
                    self._add_to_low(newElement)
            self._balance()
            
        def _add_to_low(self, newElement):
            #print 'adding the low'
            heapq.heappush(self.h_low, -1*newElement)
            self.lCount +=1
    
        def _add_to_high(self, newElement):
            #print 'adding the high'
            heapq.heappush(self.h_high, newElement)
            self.hCount +=1
            
        def _balance(self):
            """
            We only rebalance if one of the two heaps has 2 or more elements than the other
            """
            if self.lCount >= self.hCount + 2: 
                elementToMove = heapq.heappop(self.h_low) * -1 
                self.lCount -=1
                self._add_to_high(elementToMove)
            elif self.hCount >= self.lCount +2:
                elementToMove = heapq.heappop(self.h_high) 
                self.hCount -=1
                self._add_to_low(elementToMove)
            
    
        def findMedian(self):
            """
            Returns the median of current data stream
            :rtype: float
            """
            
            numElements = self.hCount + self.lCount
            if numElements %2 == 0:
                return (self.h_low[0]*(-1) + self.h_high[0])/2.0
            elif self.hCount > self.lCount:
                return self.h_high[0]
            else:
                return self.h_low[0]*-1

Log in to reply
 

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