◈ Python ◈ Ordered Set ◈ Data Stream ◈ Design
Glossary:
m: The number of most recent elements to consider.
k: The number of smallest and largest elements to exclude.
arr: The array storing all added numbers
Glossary:
- $m$: The number of most recent elements to consider.
- $k$: The number of smallest and largest elements to exclude.
- $arr$: The array storing all added numbers.
- $ss$: The sorted set storing the last $m$ elements.
- $ix$: The index of the last added element in $arr$.
- $sm$: The sum of the elements in the middle range (excluding the smallest $k$ and largest $k$ elements).
- $ix_{start}$: The starting index of the middle range in $ss$.
- $ix_{end}$: The ending index of the middle range in $ss$.
- $S$: Small element
- $M$: Middle element
- $L$: Large element
- $|ss|$: Size of the sorted set $ss$

Mathematical Intuition:

The problem requires calculating the moving average of the last $m$ numbers, excluding the $k$ smallest and $k$ largest elements. We maintain a sorted set $ss$ containing at most the last $m$ added numbers. The core idea is to efficiently update the sum $sm$ of the elements used for the average calculation.

Imagine ss divided into three sections as shown in the figure 'Three sections of sorted set ss':
```
Three sections of sorted set ss
+-----------+-----------+-----------+
|  Small  |  Middle   |  Large  |
|   k=3   |  m-2k     |   k=3   |
+-----------+-----------+-----------+
```

1. **`get_ss_ix_range_considered()`**: This function determines the indices in $ss$ that correspond to the elements used for calculating the average. Since we exclude the $k$ smallest and $k$ largest, the relevant range is from index $k$ to $|ss| - k - 1$. If the size of $ss$ is less than or equal to $2k$, there are no elements in the middle range, and the function returns `None`.

2. **`remove_ix_arr(ix_arr)`**: This function removes the element at index `ix_arr` from the considered set. It finds the corresponding element in $ss$ and removes it. The crucial part is updating $sm$. We determine the position of the removed element within $ss$ and the before and after states of the sorted set as shown in the figures 'Removing from Small (Remove S1)', 'Removing from Large (Remove L1)', and 'Removing from Middle (Remove M1)':
```
Removing from Small (Remove S1)
      0         k=3         m-k=7         m-1=9
ss:   [S][S][S]...[M][M][M]...[L][L][L]  (Before)
     ^
     |
  Removed

      0         k=3         m-k=7         m-1=9
ss:   [S][S][M]...[M][M]...[L][L][L]  (After)
sm:   (Adjust sm - a Middle element became Small)

Removing from Large (Remove L1)
      0         k=3         m-k=7         m-1=9
ss:   [S][S][S]...[M][M][M]...[L][L][L]  (Before)
                               ^
                               |
                            Removed

      0         k=3         m-k=7         m-1=9
ss:   [S][S][S]...[M][M]...[M][L][L]  (After)
sm:   (Adjust sm if the element moving into Large was previously contributing to the sum)

Removing from Middle (Remove M1)
      0         k=3         m-k=7         m-1=9
ss:   [S][S][S]...[M][M][M]...[L][L][L]  (Before)
        ^
        |
    Removed

      0         k=3         m-k=7         m-1=9
ss:   [S][S][S]...[M][M]...[L][L][L]  (After)
sm:   (Subtract M1 from sm)
```

3. **`add_ix_arr(ix_arr)`**: This function adds the element at index `ix_arr` to the considered set. It adds the element to $ss$. Then we determine the position of the newly added element within $ss$ and the before and after states of the sorted set as shown in the figures 'Adding to Small (Add S4)', 'Adding to Large (Add L4)', and 'Adding to Middle (Add M2)':
```
Adding to Small (Add S4)
      0           k      m-k         m-1
ss:   [S][S][S]...[M][M][M]...[L][L][L]  (Before)
            ^
            |
            Add
ss:   [S][S][S]...[S][M][M][M]...[L][L][L]  (After)
sm:   (Adjust sm - a Middle element became Small)

Adding to Large (Add L4)
      0           k     m-k         m-1
ss:   [S][S][S]...[M][M][M]...[L][L][L]  (Before)
                                ^
                                |
                                Add
ss:   [S][S][S]...[M][M][M][L]...[L][L][L]  (After)  
sm:   (Adjust sm - a Middle element became Large)

Adding to Middle (Add M2)
      0           k      m-k        m-1
ss:   [S][S][S]...[M][M][M]...[L][L][L]  (Before)
sm:   (Sum of elements in the Middle section)

ss:   [S][S][S]...[M][M][M][M]...[L][L][L]  (After)
sm:   (Add M2 to sm)
```

4. **`addElement(num)`**: This function adds a new number to the data structure. First, it removes the oldest element if the size of $ss$ has reached $m$. Then, it adds the new number and updates $ss$ and $sm$ using `add_ix_arr()`.

5. **`calculateMKAverage()`**: This function calculates the MKAverage. It checks if there are enough elements ($|ss| \ge m$). If not, it returns -1. Otherwise, it calculates the average by dividing $sm$ by the number of elements in the middle range ($|ss| - 2k$).

Mathematical Proof of Correctness:

The correctness hinges on the proper maintenance of $ss$ and $sm$. $ss$ always contains the last $m$ elements in sorted order. $sm$ represents the sum of the middle elements. `remove_ix_arr` and `add_ix_arr` correctly update $sm$ based on the position of the added/removed element in $ss$. `addElement` ensures that only the last $m$ elements are considered. `calculateMKAverage` uses the correctly maintained $sm$ and the size of the middle range to compute the average.

Algorithm:

1. Maintain a sorted set $ss$ of the last $m$ elements.
2. Maintain the sum $sm$ of the middle elements (excluding the $k$ smallest and $k$ largest).
3. `addElement(num)`:
    a. If $|ss| \ge m$, remove the oldest element from $ss$ and update $sm$.
    b. Add $num$ to $ss$.
    c. Update $sm$ based on the position of $num$ in $ss$.
4. `calculateMKAverage()`:
    a. If $|ss| < m$, return -1.
    b. Calculate the average: $\lfloor \frac{sm}{m - 2k} \rfloor$.

Time and Space Complexity:

Time Complexity:
- `addElement()`: $O(\log m)$ for adding/removing from the sorted set.
- `calculateMKAverage()`: $O(1)$.
Overall per operation: $O(\log m)$

Space Complexity: $O(m)$ to store the sorted set.
# Code
```python3 []
from sortedcontainers import SortedSet as SS
import math
class MKAverage:

    def __init__(self, m: int, k: int):
        self.arr = []
        self.ss = SS()
        self.ix = -1
        self.m = m 
        self.k = k
        self.sm = 0 # sum of coinsidered elements
        assert self.m > 0
        assert self.m > (self.k * 2)
        print(f"m,k {m,k}")
    def get_ss_ix_range_considered(self):
        if(
            (len(self.ss) <= 2 * self.k)
        ):
            return None, None
        ix_start = self.k
        ix_end = len(self.ss)-self.k -1
        return ix_start, ix_end
    def remove_ix_arr(self, ix_arr):
        assert 0 <= ix_arr < len(self.arr)
        item_to_remove = (self.arr[ix_arr],ix_arr)
        # update sum of considered numbers 
        # consider the last m numbers in three intervals (1-based indexing) stored in the sorted set self.ss
        # [ ... ][ k+1 ... m-k ][ ... ]
        # [  I  ][      II     ][ III ]
        # Intervals - I, III are filled first , in a balanced way
        # Only once I, III are filled with each k elements does II start getting populated
        # If any element trying to be removed from in I , I pulls leftmost element from II
        # Then any element trying to be removed in III , III pulls rightmost element from II
        # If removed number is in I -> -ss[ix_start] ( interval - I will suck up ix_start )
        # If removed number is in II -> +ss[ix_ss_new] 
        # If removed number is in III , +ss[ix_end] ( interval - |I will suck up ix_start )
        ix_ss_item_to_remove = self.ss.index(item_to_remove)
        ix_start,ix_end = self.get_ss_ix_range_considered()
        if(None in [ix_start,ix_end]):
            # there are less than 2 * k elements in self.ss so no elements in interval - II 
            pass
        elif(ix_ss_item_to_remove<ix_start):
            self.sm += -self.ss[ix_start][0]
        elif(ix_start<=ix_ss_item_to_remove<=ix_end):
            self.sm += -self.ss[ix_ss_item_to_remove][0]
        elif(ix_end<ix_ss_item_to_remove):
            self.sm += -self.ss[ix_end][0]
            pass
        self.ss.remove(item_to_remove)
    def add_ix_arr(self, ix_arr):
        assert 0 <= ix_arr < len(self.arr)
        item_to_add = (self.arr[ix_arr],ix_arr)
        # update sum of considered numbers 
        # consider the last m numbers in three intervals (1-based indexing) stored in the sorted set self.ss
        # [ ... ][ k+1 ... m-k ][ ... ]
        # [  I  ][      II     ][ III ]
        # Intervals - I, III are filled first , in a balanced way
        # Only once I, III are filled with each k elements does II start getting populated
        # Then any element trying to land in I , it pushes its rightmost element to II
        # Then any element trying to land in III , it pushes its leftmost element to II
        # If new number is in I -> +ss[ix_start] ( interval - I will push up ix_start to interval - II  )
        # If new number is in II -> +ss[ix_ss_new] 
        # If new number is in III , +ss[ix_end] ( interval - III will push up ix_end to interval - II  )
        self.ss.add(item_to_add)
        assert len(self.ss) <= self.m
        ix_ss_item_to_add = self.ss.index(item_to_add)
        ix_start,ix_end = self.get_ss_ix_range_considered()
        if(None in [ix_start,ix_end]):
            # there are less than 2 * k elements in self.ss so any new elements will move to interval I or III and there will be no change in interval - II
            pass
        elif(ix_ss_item_to_add<ix_start):
            self.sm += self.ss[ix_start][0]
        elif(ix_start<=ix_ss_item_to_add<=ix_end):
            self.sm += self.ss[ix_ss_item_to_add][0]
        elif(ix_end<ix_ss_item_to_add):
            self.sm += self.ss[ix_end][0]
            pass
    def addElement(self, num: int) -> None:
        assert len(self.ss) <= self.m
        # 1. remove index that has become stale (only consider last m elements as per P.S.)
        while(len(self.ss)  >= self.m ):
            ix_arr_to_remove = self.ix - (self.m -1)
            self.remove_ix_arr(ix_arr_to_remove)
        # 2.add new number
        self.ix +=1
        self.arr.append(num)
        self.add_ix_arr(self.ix)
        # print(f"+{num} | sm={self.sm}, ss={self.ss} , considered ixs={self.get_ss_ix_range_considered()}")

    def calculateMKAverage(self) -> int:
        if(len(self.ss) < self.m):
            return -1
        ix_start,ix_end = self.get_ss_ix_range_considered()
        if(None in [ix_start,ix_end]):
            return -1
        avg = math.floor(self.sm / (ix_end +1 - ix_start))
        return avg
        



# Your MKAverage object will be instantiated and called as such:
# obj = MKAverage(m, k)
# obj.addElement(num)
# param_2 = obj.calculateMKAverage()
```
['Design', 'Data Stream', 'Ordered Set', 'Python3']