average.py

class MKAverage:

    def __init__(self, m: int, k: int):
        self.q = SkipList(6, 0.5)
        self.m = m
        self.k = k
        
        self.counter = 0
        
        self.sum_ = 0

    def addElement(self, new: int) -> None:
        # Hacks
        q = self.q

        k = self.k
        m = self.m
        
        # If list not full, just append
        if self.counter < m:
            q.push(new)
            self.counter += 1
            
            if self.counter == 2 * k:
                self.top_k = q[k]
                self.bottom_k = q[k - 1]
            
            elif self.counter > 2 * k:
                if new >= self.top_k.key:
                    self.sum_ += self.top_k.key
                    self.top_k = self.top_k.forward[0]
                elif new <= self.bottom_k.key:
                    self.sum_ += self.bottom_k.key
                    self.bottom_k = self.bottom_k.backward[0]
                else:
                    self.sum_ += new
            return
        
        # Remove node from the left
        node = q.popleft()
        if node.key >= self.top_k.key:
            self.top_k = self.top_k.backward[0]
            self.sum_ -= self.top_k.key
        elif node.key <= self.bottom_k.key:
            self.bottom_k = self.bottom_k.forward[0]
            self.sum_ -= self.bottom_k.key
        else:
            self.sum_ -= node.key
            
        # Insert new node
        q.push(new)
        if new >= self.top_k.key:
            self.sum_ += self.top_k.key
            self.top_k = self.top_k.forward[0]
        elif new <= self.bottom_k.key:
            self.sum_ += self.bottom_k.key
            self.bottom_k = self.bottom_k.backward[0]
        else:
            self.sum_ += new
            

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