1825. Finding MK Average (Hard) - TengnanYao/daily_leetcode GitHub Wiki

from sortedcontainers import SortedList
class MKAverage(object):

    def __init__(self, m, k):
        """
        :type m: int
        :type k: int
        """
        self.m = m
        self.k = k
        self.s = []
        self.q = []
        self.sum = 0

    def addElement(self, num):
        """
        :type num: int
        :rtype: None
        """
        self.q.append(num)
        if len(self.s) == self.m:
            i = bisect.bisect_left(self.s, num)
            if self.k <= i <= self.m - self.k:
                self.sum += num
            elif i < self.k:
                self.sum += self.s[self.k - 1]
            else:
                self.sum += self.s[-self.k]
            self.s.add(num)
            val = self.q.pop(0)
            j = bisect.bisect_left(self.s, val)
            if self.k <= j <= self.m - self.k:
                self.sum -= val
            elif j < self.k:
                self.sum -= self.s[self.k]
            else:
                self.sum -= self.s[-self.k - 1]
            self.s.remove(val)
        if len(self.q) == self.m and self.sum == 0:
            self.s = SortedList(self.q)
            self.sum = sum(self.s[self.k : -self.k])

    def calculateMKAverage(self):
        """
        :rtype: int
        """
        if len(self.q) < self.m:
            return -1
        return int(round(self.sum / (self.m - 2 * self.k)))


# Your MKAverage object will be instantiated and called as such:
# obj = MKAverage(m, k)
# obj.addElement(num)
# param_2 = obj.calculateMKAverage()