Python O(nlogK) solution using two heaps


  • 0
    J

    Inspired by this post:
    https://discuss.leetcode.com/topic/74679/o-n-log-n-time-c-solution-using-two-heaps-and-a-hash-table

    Use two heaps, one for smaller elements, one for larger elements. Each element in the heap contains the value and index. The main problem of using heapq in Python is that removal costs O(n) time. When we reach to index i, we need to remove those elements whose indices are smaller or equals to i - k (we could call them useless elements). My method is to leave those useless elements in the heap if they are not on the top of the heap so that they wouldn't affect the answer. Also I used a variable self.balance to denote where those useless elements are. If the useless elements are on the larger heap (check by comparing value with the top one on the larger heap), self.balance plus 1, (minus 1 if in the smaller heap). So the actually useful elements in larger heap will be len(l) - self.balance. Then we compare it with len(s) to determine where we should add the next element and we should always add the new element into the shorter heap. If the useless elements are on the top of heaps, we just pop them and change the balance until the one on the top is inside the right range (index is bigger than i - k). We could get the median by check top elements on both heap.

    '''

    def medianSlidingWindow(self, nums, k):
        def addNumber(num, i):
            if len(l) - len(s) > self.balance:
                val, idx = heapq.heappushpop(l, (num, i))
                heapq.heappush(s, (-val, idx))
            else:
                val, idx = heapq.heappushpop(s, (-num, i))
                heapq.heappush(l, (-val, idx))
    
    
        def removeNumber(idx):
            while l and l[0][1] <= idx:
                heapq.heappop(l)
                self.balance -= 1
    
            while s and s[0][1] <= idx:
                heapq.heappop(s)
                self.balance += 1
    
    
        s, l, self.balance = [], [], 0
        res = []
        for i in xrange(len(nums)):
            if i < k:
                addNumber(nums[i], i)
            else:
                if nums[i - k] > l[0][0]:
                    self.balance += 1
                elif nums[i - k] < l[0][0]:
                    self.balance -= 1
                else:
                    for val, idx in l:
                        if idx == i - k:
                            self.balance += 1
                            break
                    else:
                        self.balance -= 1
                addNumber(nums[i], i)
                removeNumber(i - k)
            if i >= k - 1:
                median = l[0][0] * 1.0 if k % 2 == 1 else (l[0][0] - s[0][0]) / 2.0
                res.append(median)
        return res
    

    '''


  • 0
    L

    Why can't you do removenumber before addnumber?


Log in to reply
 

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