Finding MK Average LeetCode solution in Python

1
from sortedcontainers import SortedList
2

3

4
class MKAverage:
5

6
    MAX_NUM = 10**5
7

8
    def __init__(self, m: int, k: int):
9

10
        self.m = m
11
        self.k = k
12

13
        # sorted list
14
        self.sl = SortedList([0] * m)
15
        # sum of k smallest elements
16
        self.sum_k = 0
17
        # sum of m - k smallest elements
18
        self.sum_m_k = 0
19

20
        # queue for the last M elements if the stream
21
        self.q = deque([0] * m)
22

23
    def addElement(self, num: int) -> None:
24
        # Time: O(logm)
25

26
        m, k, q, sl = self.m, self.k, self.q, self.sl
27

28
        # update q
29
        q.append(num)
30
        old = q.popleft()
31

32
        # remove the old num
33
        r = sl.bisect_right(old)
34
        # maintain sum_k
35
        if r <= k:
36
            self.sum_k -= old
37
            self.sum_k += sl[k]
38
        # maintain sum_m_k
39
        if r <= m - k:
40
            self.sum_m_k -= old
41
            self.sum_m_k += sl[m - k]
42
        # remove the old num
43
        sl.remove(old)
44

45
        # add the new num
46
        r = sl.bisect_right(num)
47
        if r < k:
48
            self.sum_k -= sl[k - 1]
49
            self.sum_k += num
50
        if r < m - k:
51
            self.sum_m_k -= sl[m - k - 1]
52
            self.sum_m_k += num
53

54
        sl.add(num)
55

56
        return
57

58
    def calculateMKAverage(self) -> int:
59
        # Time: O(1)
60
        if self.sl[0] == 0:
61
            return -1
62
        return (self.sum_m_k - self.sum_k) // (self.m - self.k * 2)