Python solution O(n log k) with only two heaps and explanation


  • 0
    C

    Coding style might not be good, hope you like it.

    class Solution(object):
        def medianSlidingWindow(self, nums, k):
        
        from heapq import heappush, heappop, heapify
        
        lMaxHeap = []
        # for every pair in heap, not only record the key, but also its position in nums
        firstList = [(nums[i], i) for i in range(k)]
        heapify(firstList)
        rMinHeap = firstList
        if k == 1:
            return [float(i) for i in nums]
        # initialize two heap: size of Right min heap  - size of Left max heap >= 1
        # in the following, abbreviate right min heap as right; left max heap as left
        # when k is odd, left size == right size - 1;
        # when k is even, left size == right size
        while len(rMinHeap) - len(lMaxHeap) >= 2:
            pop = heappop(rMinHeap)
            # implement max heap by key * (-1) in min heap
            heappush(lMaxHeap, ( -1 * pop[0], pop[1]))
        
        medianList = [float(rMinHeap[0][0]) if k%2 == 1 else float(rMinHeap[0][0] - lMaxHeap[0][0])/2]
        
        for i in xrange(k, len(nums)):
            # if balance >= 2 -> push(left , pop(right))
            # elif balance <= 2 -> push(left, pop(right))
            balance = 0
            # add cases: when added num greater or equal to right top of min heap, add it to right. \
            # Otherwise, add it to left
            if nums[i] >= rMinHeap[0][0]:
                balance += 1
                heappush(rMinHeap, (nums[i], i))
            else:
                balance -= 1
                heappush(lMaxHeap, (-nums[i], i))
            # remove case:
            # (1) the removing element is in top of either list -> remove until top is in current list
            # at the same time, change balance 
            # (2) the removing element is not in both top of heap -> just change balance
                
            if lMaxHeap[0][1] == i-k:
                heappop(lMaxHeap)
                while (lMaxHeap) and (lMaxHeap[0][1] <= i - k):
                    heappop(lMaxHeap) 
                balance += 1
            elif rMinHeap[0][1] == i - k:
                heappop(rMinHeap)
                while (rMinHeap) and (rMinHeap[0][1] <= i - k):
                    heappop(rMinHeap)
                balance -= 1
            elif rMinHeap[0][0] >= nums[i - k]:
                balance += 1
            else:
                balance -= 1
            
            #balance num in both heap
            if balance > 0:
                popElement = heappop(rMinHeap)
                heappush(lMaxHeap, (-popElement[0], popElement[1]))
                while (rMinHeap) and (rMinHeap[0][1] <= i - k):
                    heappop(rMinHeap)
            elif balance < 0:
                popElement = heappop(lMaxHeap)
                heappush(rMinHeap,(-popElement[0], popElement[1]))
                while (lMaxHeap) and (lMaxHeap[0][1] <= i - k):
                    heappop(lMaxHeap)  
    
            medianList.append(float(rMinHeap[0][0]) if k%2 == 1 else float(rMinHeap[0][0] - lMaxHeap[0][0])/2)
    
                    
        return medianList

Log in to reply
 

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