Problem statement

https://leetcode.com/problems/finding-mk-average/

Solution

The idea is similar to problem 480. Sliding Window Median, here we also have sliding window, but we need to find not median, but some other statistics: sum of some middle elements in sorted window. Imagine for the purpose of example that m = 7 and k = 2. Then we have widow of size 7: a1 < a2 < a3 < a4 < a5 < a6 < a7 and we need to evaluate the sum of 3 numbers a3 + a4 + a5. Let us solve two problems: evaluate a1 + a2 + a3 + a4 + a5 and then substitute a1 + a2. How we are going to evaluate sum a1 + a2 for example?

The idea is to keep 2 heaps: one for smallest 2 numbers in sliding window and another is for the 5 biggest numbers. We do it in very similar way as problem 480, however we need to also evaluate sums. The trick is to use lazy deletion from heap, that is sometimes it happen, that there are some old elements in our window we need to be deleted later. If you know what is lazy deletion from heap, I think you got the idea, if not, please google it. The difficult part is how to update sum in heap as well, we need to do it carefully, because we have lazy updates, we need to make sure when we add or delete element, that it is element which will not be lazy update.

Finally, we need to solve two separate problems: one for parts with sizes 2 vs 5 and another with sizes 5 vs 2. Also I use trick where I fill heap with zeroes at first, so no need to check case when we do not have enough elements it parts.

Complexity

Time complexity is O(n log n), where n is total number of calls of addElement. Space complexity is O(n).

class MKAverage:
    def __init__(self, m: int, k: int):
        self.m, self.k = m, k
        self.arr = [0]*m
        self.lh1, self.rh1 = self.heap_init(m, k)
        self.lh2, self.rh2 = self.heap_init(m, m - k)
        self.score = 0
        
    def heap_init(self, p1, p2):
        h1 = [(0, i) for i in range(p1-p2, p1)]
        h2 = [(0, i) for i in range(p1-p2)]
        heapq.heapify(h1)
        heapq.heapify(h2)
        return (h1, h2)
        
    def update(self, lh, rh, m, k, num):
        score, T = 0, len(self.arr)
        if num > rh[0][0]:
            heappush(rh, (num, T))        
            if self.arr[T - m] <= rh[0][0]:
                if rh[0][1] >= T - m: score += rh[0][0]
                score -= self.arr[T - m]
                heappush(lh, (-rh[0][0], rh[0][1]))
                heappop(rh)
        else:
            heappush(lh, (-num, T))       
            score += num
            if self.arr[T - m] >= rh[0][0]: 
                heappush(rh, (-lh[0][0], lh[0][1]))
                q = heappop(lh)
                score += q[0]
            else:
                score -= self.arr[T - m]

        while lh and lh[0][1] <= T - m: heappop(lh)  # lazy-deletion
        while rh and rh[0][1] <= T - m: heappop(rh)  # lazy-deletion
        return score
        
    def addElement(self, num):
        t1 = self.update(self.lh1, self.rh1, self.m, self.k, num)
        t2 = self.update(self.lh2, self.rh2, self.m, self.m - self.k, num)
        self.arr.append(num)
        self.score += (t2 - t1)
    
    def calculateMKAverage(self):
        if len(self.arr) < 2*self.m: return -1
        return self.score//(self.m - 2*self.k)

If you have any questions, feel free to ask. If you like solution and explanations, please Upvote!