1825. Finding MK Average
On This Page
from sortedcontainers import SortedList
class MKAverage:
def __init__(self, m: int, k: int):
self.m = m
self.k = k
self.sl = SortedList()
self.q = deque()
self.s = 0
def addElement(self, num: int) -> None:
self.q.append(num)
if len(self.q) == self.m:
self.sl = SortedList(self.q)
self.s = sum(self.sl[self.k : -self.k])
elif len(self.q) > self.m:
i = self.sl.bisect_left(num)
if i < self.k:
self.s += self.sl[self.k - 1]
elif self.k <= i <= self.m - self.k:
self.s += num
else:
self.s += self.sl[self.m - self.k]
self.sl.add(num)
x = self.q.popleft()
i = self.sl.bisect_left(x)
if i < self.k:
self.s -= self.sl[self.k]
elif self.k <= i <= self.m - self.k:
self.s -= x
else:
self.s -= self.sl[self.m - self.k]
self.sl.remove(x)
def calculateMKAverage(self) -> int:
return -1 if len(self.sl) < self.m else self.s // (self.m - self.k * 2)
# Your MKAverage object will be instantiated and called as such:
# obj = MKAverage(m, k)
# obj.addElement(num)
# param_2 = obj.calculateMKAverage()