Python O(n*log(k)) using 2 heaps and a set with explanation (AC: 462 ms)


  • 0
    P

    The code (below) is a bit verbose and it took me a lot of time to implement it because of silly errors but finally got it Accepted (~462ms). Even though the code looks verbose if the bigger picture is clear, it should seem simple. The idea is similar to what is pointed out in other top solutions i.e. maintaining 2 heaps, left and right to store max of left half and min of right half of the sorted window, without actually sorting the window. The following condition must hold, right.length-left.length <= 1 for the median to be sandwiched between these elements.

    The key challenge is that. whenever the window shifts, a new element should get added to left/right heap and one old element that just got left out of the window must be removed from the left/right heap. While adding is simple O(log k), removing the element that got left out of the window from the heap efficiently is a bit tricky. That's because Python heaps are stored as arrays and if we were to remove an element and re-heapify the array then it would cost O(k), making the overall solution O(nk) instead of O(n*log k) and likely get a TLE with Python. Instead, I marked removed elements and maintained them in a set without actually removing the elements from the heap. Someone even hinted at this idea in Python heapq module docs in context of implementation of schedulers. This allowed me to do O(1) removal from the heap at the cost of some memory, because the removed set grows and removal of elements from it is handled lazily.

    Using this idea now removal becomes O(1). The add O(log k) step dominates. The window has to be shifted about (N-k) times. So overall ~ (N-k) * O(log k) ops are needed, making the solution O(n*log(k)) overall. Code with some inline comments below:

    import heapq as hq
    minf, inf = float('-inf'), float('inf')
    
    class MinHeap():
        def __init__(self):
            self.h = []
            self.removed = set()
        
        def push(self, v):
            hq.heappush(self.h, v)
        
        def pop(self):
            while True:
                v = hq.heappop(self.h)
                if v in self.removed:
                    self.removed.remove(v)
                else:
                    break
            return v
    
        def top(self):
            while True:
                v = self.h[0]
                if v not in self.removed:
                    return v
                else:
                    v = hq.heappop(self.h)
                    self.removed.remove(v)
    
        def remove(self, v):
            # Assumes called with a value v, that is already there in the heap
            if v not in self.removed:
                self.removed.add(v)
        
        @property
        def length(self):
            return len(self.h) - len(self.removed)
    
    class MaxHeap():
        def __init__(self):
            self.min_heap = MinHeap()
        
        def negative(self, v):
            return tuple(map(lambda x: -x, v))
    
        def push(self, v):
            self.min_heap.push(self.negative(v))
        
        def pop(self):
            v = self.min_heap.pop()
            return self.negative(v)
        
        def top(self):
            v = self.min_heap.top()
            return self.negative(v)
    
        def remove(self, v):
            # Assumes only those elements shall be removed which are already in the heap
            # The caller should know that this method is being called with a value that exists in the heap
            self.min_heap.remove(self.negative(v))
    
        @property
        def length(self):
            return self.min_heap.length
    
    
    def add_and_rebalance(n, left, right):
        """
        Adds an element to the left/right heap and rebalances such that, right (min heap)
        is either equal or just one unit more than left (max heap).
        """
        if n >= right.top():
            right.push(n)
        else:
            left.push(n)
    
        if left.length > right.length:
            right.push(left.pop())
        
        if right.length > (left.length + 1):
            left.push(right.pop())
    
    
    def remove_and_rebalance(n, left, right):
        """
        Removes an element from the left/right heap and rebalances such that, right (min heap)
        is either equal or just one unit more than left (max heap).
        """
        if n >= right.top():
            right.remove(n)
        else:
            left.remove(n)
    
        if left.length > right.length:
            right.push(left.pop())
        
        if right.length > (left.length + 1):
            left.push(right.pop())
    
    
    class Solution(object):
        def medianSlidingWindow(self, nums, k):
            """
            :type nums: List[int]
            :type k: int
            :rtype: List[float]
            """
            L = len(nums)
            if k == 1:
                return map(float, nums)
            
            left, right = MaxHeap(), MinHeap()
            left.push((minf, -1))
            right.push((inf, -1))
            
            for i in range(k):
                n = nums[i]
                add_and_rebalance((n,i), left, right)
    
            medians = []
            for i in range(k, L+1):
        
                if right.length == left.length:
                    ltop = left.top()[0]
                    rtop = right.top()[0]
                    m = float(ltop + rtop) / 2
                elif right.length == left.length + 1:
                    rtop = right.top()[0]
                    m = float(rtop)
    
                medians.append(m)
                
                if i == L:
                    break
    
                # Remove element left out of window
                rm_index = i-k
                rm_num = nums[rm_index]
    
                remove_and_rebalance((rm_num, rm_index), left, right)
                add_and_rebalance((nums[i], i), left, right)
    
            return medians
    

Log in to reply
 

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