mkaverage.py

import array
import collections
import bisect

class MKAverage:

    def __init__(self, m: int, k: int):
        # self.q = collections.deque()
        self.q = array.array("L")
        self.p = None
        self.m = m
        self.k = k
        
        self.sum_ = None
        self.qi = 0

    def addElement(self, new: int) -> None:
        # Hacks
        q = self.q
        p = self.p
        k = self.k
        m = self.m
        
        # If list not full, just append
        if len(q) < m:
            q.append(new)
            
            if len(q) == m:
                p = sorted(q)
                self.p = collections.deque(p)
                self.sum_ = sum(p[k:-k])
            return
            
        # Update sorted list
        old = q[self.qi]
        q[self.qi] = new
        self.qi = (self.qi + 1) % m
        
        # Update removal of old numbers
        i = bisect.bisect_left(p, old)
        if i < k:
            self.sum_ -= p[k]
        elif i >= m - k:
            self.sum_ -= p[-k - 1]
        else:
            self.sum_ -= old
        del p[i]
        
        # Update insertion of new number
        j = bisect.bisect_left(p, new)
        if j < k:
            self.sum_ += p[k - 1]
        elif j >= m - k:
            self.sum_ += p[-k]
        else:
            self.sum_ += new
        p.insert(j, new)
        

    def calculateMKAverage(self) -> int:
        if len(self.q) < self.m:
            return -1
        return self.sum_ // (self.m - 2 * self.k)