# Finding MK Average
**Difficulty:** HARD
[External](https://leetcode.com/problems/finding-mk-average)
Canonical: https://scaleengineer.com/dsa/problems/finding-mk-average
**Patterns:** [Design](https://scaleengineer.com/dsa/patterns/design), [Data Stream](https://scaleengineer.com/dsa/patterns/data-stream)
**Data structures:** Heap (Priority Queue), Ordered Set, Queue
---
## Problem
You are given two integers, `m` and `k`, and a stream of integers. You are tasked to implement a data structure that calculates the **MKAverage** for the stream.

The **MKAverage** can be calculated using these steps:

1. If the number of the elements in the stream is less than `m` you should consider the **MKAverage** to be `-1`. Otherwise, copy the last `m` elements of the stream to a separate container.
2. Remove the smallest `k` elements and the largest `k` elements from the container.
3. Calculate the average value for the rest of the elements **rounded down to the nearest integer**.

Implement the `MKAverage` class:

* `MKAverage(int m, int k)` Initializes the **MKAverage** object with an empty stream and the two integers `m` and `k`.
* `void addElement(int num)` Inserts a new element `num` into the stream.
* `int calculateMKAverage()` Calculates and returns the **MKAverage** for the current stream **rounded down to the nearest integer**.

**Example 1:**

**Input**
["MKAverage", "addElement", "addElement", "calculateMKAverage", "addElement", "calculateMKAverage", "addElement", "addElement", "addElement", "calculateMKAverage"]
[[3, 1], [3], [1], [], [10], [], [5], [5], [5], []]
**Output**
[null, null, null, -1, null, 3, null, null, null, 5]

**Explanation**
`MKAverage obj = new MKAverage(3, 1); 
obj.addElement(3);        // current elements are [3]
obj.addElement(1);        // current elements are [3,1]
obj.calculateMKAverage(); // return -1, because m = 3 and only 2 elements exist.
obj.addElement(10);       // current elements are [3,1,10]
obj.calculateMKAverage(); // The last 3 elements are [3,1,10].
                          // After removing smallest and largest 1 element the container will be [3].
                          // The average of [3] equals 3/1 = 3, return 3
obj.addElement(5);        // current elements are [3,1,10,5]
obj.addElement(5);        // current elements are [3,1,10,5,5]
obj.addElement(5);        // current elements are [3,1,10,5,5,5]
obj.calculateMKAverage(); // The last 3 elements are [5,5,5].
                          // After removing smallest and largest 1 element the container will be [5].
                          // The average of [5] equals 5/1 = 5, return 5
`

**Constraints:**

* `3 <= m <= 105`
* `1 < k*2 < m`
* `1 <= num <= 105`
* At most `105` calls will be made to `addElement` and `calculateMKAverage`.

# Approaches
## Brute Force with Sorting
This approach directly simulates the process described in the problem. It maintains a sliding window of the last `m` elements using a queue. For each `calculateMKAverage` call, it copies these `m` elements into a temporary list, sorts the list, and then calculates the average of the middle elements. While simple to understand, its performance suffers due to the repeated sorting.
**Time:** `addElement`: O(1)
Adding and removing from a queue are constant time operations.

`calculateMKAverage`: O(m log m)
The dominant operation is sorting the list of `m` elements. · **Space:** O(m)
We need a queue to store the last `m` elements. The temporary list in `calculateMKAverage` also takes `O(m)` space.
**Pros:** Easy to understand and implement.; The logic directly follows the problem statement.
**Cons:** The `calculateMKAverage` method is inefficient, with a time complexity of `O(m log m)`.; For streams with frequent average calculations, this approach can be very slow, potentially leading to a 'Time Limit Exceeded' error in a competitive programming context.
### Explanation
The core idea is to use a `Queue` data structure to keep track of the last `m` elements added to the stream. This naturally handles the sliding window requirement.

When `addElement(num)` is called, the new number is added to the queue. To ensure the queue only contains the last `m` elements, we check its size after adding. If it exceeds `m`, we remove the oldest element, which is at the front of the queue.

When `calculateMKAverage()` is called, the main work begins. First, a size check ensures we have enough elements. If not, `-1` is returned. Otherwise, the current `m` elements from the queue are copied into a new list. This list is then sorted, which is the most time-consuming step (`O(m log m)`). After sorting, the smallest `k` elements are at the beginning of the list, and the largest `k` are at the end. We can simply ignore them by summing the elements from index `k` up to (but not including) index `m - k`. The final average is this sum divided by the number of elements considered (`m - 2*k`).

```java
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;

class MKAverage {
    private int m, k;
    private Queue<Integer> stream;

    public MKAverage(int m, int k) {
        this.m = m;
        this.k = k;
        this.stream = new LinkedList<>();
    }
    
    public void addElement(int num) {
        stream.offer(num);
        if (stream.size() > m) {
            stream.poll();
        }
    }
    
    public int calculateMKAverage() {
        if (stream.size() < m) {
            return -1;
        }

        List<Integer> window = new ArrayList<>(stream);
        Collections.sort(window);

        long sum = 0;
        for (int i = k; i < m - k; i++) {
            sum += window.get(i);
        }

        return (int) (sum / (m - 2 * k));
    }
}
```
### Algorithm
- **Initialization**: In the constructor, initialize a `Queue` (e.g., `LinkedList`) to store the stream elements, and store the values of `m` and `k`.
- **`addElement(num)`**: 
  - Add the new element `num` to the tail of the queue.
  - If the size of the queue is greater than `m`, remove the element from the head of the queue (`poll()`). This maintains a sliding window of the last `m` elements.
- **`calculateMKAverage()`**: 
  - Check if the number of elements in the queue is less than `m`. If it is, return `-1`.
  - Create a temporary `ArrayList` and copy all elements from the queue into it.
  - Sort the temporary list using a standard sorting algorithm (e.g., `Collections.sort()`).
  - Initialize a variable `sum` to zero.
  - Iterate through the sorted list from index `k` to `m - 1 - k` (inclusive) and add each element to `sum`.
  - Calculate the average by dividing `sum` by the count of middle elements, which is `m - 2*k`.
  - Return the result, which is automatically floored due to integer division.

## Optimization with a Single Sorted Map
This approach improves upon the brute-force method by avoiding the expensive sorting step in every `calculateMKAverage` call. It maintains the sliding window of `m` elements in a sorted state using a `TreeMap`, which stores numbers and their frequencies. While adding elements is efficient (`O(log m)`), calculating the average still requires iterating through the map's elements, resulting in `O(m)` complexity for that step.
**Time:** `addElement`: O(log m)
`TreeMap` operations (put, get, remove) take logarithmic time in the size of the map, which is at most `m`.

`calculateMKAverage`: O(m)
In the worst case, we might have to iterate through all `m` elements if they are all distinct. · **Space:** O(m)
Space is required for the queue (`O(m)`) and the `TreeMap` (at most `m` distinct elements, so `O(m)`).
**Pros:** More efficient than the brute-force approach.; `addElement` is fast, with a logarithmic time complexity.
**Cons:** The `calculateMKAverage` method still requires iterating through the elements in the sorted map, leading to a time complexity of `O(m)` in the worst case (when all elements are distinct).; While better than brute force, it may still be too slow if `m` is very large.
### Explanation
To avoid re-sorting `m` elements every time, we can use a data structure that maintains sorted order as elements are added and removed. A `TreeMap` is perfect for this, as it stores key-value pairs sorted by key. We can use the numbers as keys and their frequencies as values.

We still use a `Queue` to know which element is the oldest and needs to be removed from the window. When `addElement(num)` is called, we add `num` to the queue and increment its count in the `TreeMap`. When an element is removed from the queue, its count in the `TreeMap` is decremented.

When `calculateMKAverage()` is called, we no longer need to sort. We can iterate through the `TreeMap`, which provides the elements in sorted order. We traverse the map, keeping track of how many elements we've seen. We skip the first `k`, sum the next `m - 2*k`, and then stop. This traversal takes time proportional to the number of distinct elements in the window, which is at most `m`.

```java
import java.util.LinkedList;
import java.util.Queue;
import java.util.TreeMap;
import java.util.Map;

class MKAverage {
    private int m, k;
    private Queue<Integer> window;
    private TreeMap<Integer, Integer> sortedCounts;

    public MKAverage(int m, int k) {
        this.m = m;
        this.k = k;
        this.window = new LinkedList<>();
        this.sortedCounts = new TreeMap<>();
    }

    public void addElement(int num) {
        window.offer(num);
        sortedCounts.put(num, sortedCounts.getOrDefault(num, 0) + 1);

        if (window.size() > m) {
            int oldNum = window.poll();
            sortedCounts.put(oldNum, sortedCounts.get(oldNum) - 1);
            if (sortedCounts.get(oldNum) == 0) {
                sortedCounts.remove(oldNum);
            }
        }
    }

    public int calculateMKAverage() {
        if (window.size() < m) {
            return -1;
        }

        long sum = 0;
        int count = 0;
        int elementsToSkip = k;
        int elementsToSum = m - 2 * k;

        for (Map.Entry<Integer, Integer> entry : sortedCounts.entrySet()) {
            int num = entry.getKey();
            int freq = entry.getValue();

            if (elementsToSkip > 0) {
                int canSkip = Math.min(elementsToSkip, freq);
                elementsToSkip -= canSkip;
                freq -= canSkip;
            }

            if (freq > 0 && elementsToSum > 0) {
                int canSum = Math.min(elementsToSum, freq);
                sum += (long) num * canSum;
                elementsToSum -= canSum;
            }

            if (elementsToSum == 0) {
                break;
            }
        }
        return (int) (sum / (m - 2 * k));
    }
}
```
### Algorithm
- **Initialization**: Use a `Queue` for the sliding window and a `TreeMap` to store the frequency of numbers within the window. The `TreeMap` naturally keeps the numbers sorted.
- **`addElement(num)`**: 
  - Add `num` to the queue and update its count in the `TreeMap`. This is an `O(log m)` operation.
  - If the queue size exceeds `m`, remove the oldest element (`oldNum`) from the queue. Decrement `oldNum`'s count in the `TreeMap`, removing it entirely if the count becomes zero. This is also `O(log m)`.
- **`calculateMKAverage()`**: 
  - If the queue size is less than `m`, return `-1`.
  - Iterate through the `TreeMap`'s entries (which are sorted by key).
  - Use a counter to skip the first `k` elements (considering their frequencies).
  - Sum the next `m - 2*k` elements.
  - Stop iterating once the required number of middle elements have been summed.
  - Return the sum divided by `m - 2*k`.

## Optimal Approach with Three Sorted Maps
This optimal approach achieves `O(1)` time for `calculateMKAverage` by continuously maintaining the sum of the middle elements. It partitions the `m` elements in the sliding window into three groups using three `TreeMap`s: the `k` smallest (`low`), the `m - 2k` middle (`mid`), and the `k` largest (`high`). The `addElement` operation involves adding the new element, removing the old one, and then rebalancing the three partitions, which takes `O(log m)` time. The average calculation then becomes a simple lookup.
**Time:** `addElement`: O(log m)
Each addition/removal involves a few `TreeMap` operations, each taking `O(log m)` or `O(log k)` time. The rebalancing moves a constant number of elements.

`calculateMKAverage`: O(1)
The result is pre-computed and stored in `midSum`, so retrieval is constant time. · **Space:** O(m)
Space is needed for the queue (`O(m)`) and the three `TreeMap`s, which together store `m` elements.
**Pros:** Extremely fast `calculateMKAverage` calls, which are `O(1)`.; The most efficient solution for the given problem constraints, especially when average calculations are frequent.
**Cons:** The implementation is significantly more complex than the other approaches.; The rebalancing logic requires careful handling of many edge cases to ensure correctness.
### Explanation
The key insight for the most optimal solution is to realize that the sum needed for the average can be maintained incrementally. We can split the `m` elements into three logical, sorted groups.

- `low`: A sorted container for the `k` smallest elements.
- `mid`: A sorted container for the `m - 2k` middle elements.
- `high`: A sorted container for the `k` largest elements.

We use `TreeMap`s to implement these sorted containers, storing element frequencies. We also maintain a `long midSum` to hold the sum of elements currently in the `mid` container.

When `addElement(num)` is called, we add `num` and remove the oldest element from the window. This triggers a rebalancing act. An element leaving the window is removed from its respective `TreeMap`. A new element is added. Then, we shuffle elements between the three `TreeMap`s to restore their required sizes (`k`, `m-2k`, `k`). For example, if `low` has more than `k` elements, its largest element is moved to `mid`. If `mid` has too many, its largest is moved to `high`. This process is reversed if a partition is too small. Crucially, `midSum` is updated with every element that moves into or out of the `mid` container.

With this intricate bookkeeping, the `calculateMKAverage()` method becomes trivial. It just returns the pre-calculated `midSum` divided by `m - 2*k`.

```java
import java.util.LinkedList;
import java.util.Queue;
import java.util.TreeMap;

class MKAverage {
    private int m, k;
    private Queue<Integer> q;
    private TreeMap<Integer, Integer> low, mid, high;
    private int lowSize, highSize;
    private long midSum;

    public MKAverage(int m, int k) {
        this.m = m;
        this.k = k;
        this.q = new LinkedList<>();
        this.low = new TreeMap<>();
        this.mid = new TreeMap<>();
        this.high = new TreeMap<>();
        this.lowSize = 0;
        this.highSize = 0;
        this.midSum = 0;
    }

    public void addElement(int num) {
        q.offer(num);
        int oldNum = -1;
        if (q.size() > m) {
            oldNum = q.poll();
        }

        // Remove old element if window was full
        if (oldNum != -1) {
            if (low.containsKey(oldNum)) {
                remove(low, oldNum); lowSize--;
            } else if (mid.containsKey(oldNum)) {
                remove(mid, oldNum); midSum -= oldNum;
            } else {
                remove(high, oldNum); highSize--;
            }
        }

        // Add new element
        if (!low.isEmpty() && num <= low.lastKey()) {
            add(low, num); lowSize++;
        } else if (!high.isEmpty() && num >= high.firstKey()) {
            add(high, num); highSize++;
        } else {
            add(mid, num); midSum += num;
        }

        // Rebalance partitions
        while (lowSize > k) {
            int val = low.lastKey();
            remove(low, val); lowSize--;
            add(mid, val); midSum += val;
        }
        while (highSize > k) {
            int val = high.firstKey();
            remove(high, val); highSize--;
            add(mid, val); midSum += val;
        }
        while (lowSize < k && !mid.isEmpty()) {
            int val = mid.firstKey();
            remove(mid, val); midSum -= val;
            add(low, val); lowSize++;
        }
        while (highSize < k && !mid.isEmpty()) {
            int val = mid.lastKey();
            remove(mid, val); midSum -= val;
            add(high, val); highSize++;
        }
    }

    public int calculateMKAverage() {
        if (q.size() < m) {
            return -1;
        }
        return (int) (midSum / (m - 2 * k));
    }

    private void add(TreeMap<Integer, Integer> map, int val) {
        map.put(val, map.getOrDefault(val, 0) + 1);
    }

    private boolean remove(TreeMap<Integer, Integer> map, int val) {
        if (!map.containsKey(val)) return false;
        map.put(val, map.get(val) - 1);
        if (map.get(val) == 0) {
            map.remove(val);
        }
        return true;
    }
}
```
### Algorithm
- **Initialization**: Maintain a `Queue` for the sliding window, and three `TreeMap`s: `low` for the smallest `k` elements, `mid` for the middle `m-2k` elements, and `high` for the largest `k` elements. Also, keep track of the sizes of these partitions and the sum of elements in `mid` (`midSum`).
- **`addElement(num)`**: This is the main logic.
  1.  An incoming element `num` is added, and if the window is full, an `oldNum` is removed.
  2.  First, handle the removal of `oldNum` from whichever partition (`low`, `mid`, or `high`) it resides in. Update the corresponding size and `midSum` if it was in `mid`.
  3.  Next, add the new element `num` to one of the partitions. A simple initial placement is to add it to `low` and then rebalance, or to place it based on the current boundaries (`low.lastKey()` and `high.firstKey()`).
  4.  **Rebalance**: After the addition and removal, the partition sizes might be incorrect. Move elements between adjacent partitions (`low` <-> `mid`, `mid` <-> `high`) to restore the target sizes (`k`, `m-2k`, `k`). Update `midSum` whenever an element moves into or out of the `mid` partition. This rebalancing involves a constant number of moves.
- **`calculateMKAverage()`**: 
  - If the queue size is less than `m`, return `-1`.
  - Since `midSum` is always kept up-to-date, simply return `(int) (midSum / (m - 2 * k))`. This is an `O(1)` operation.

# Solutions
### Java

```java
class MKAverage { private int m , k ; private long s ; private int size1 , size3 ; private Deque < Integer > q = new ArrayDeque <>(); private TreeMap < Integer , Integer > lo = new TreeMap <>(); private TreeMap < Integer , Integer > mid = new TreeMap <>(); private TreeMap < Integer , Integer > hi = new TreeMap <>(); public MKAverage ( int m , int k ) { this . m = m ; this . k = k ; } public void addElement ( int num ) { if ( lo . isEmpty () || num <= lo . lastKey ()) { lo . merge ( num , 1 , Integer: : sum ); ++ size1 ; } else if ( hi . isEmpty () || num >= hi . firstKey ()) { hi . merge ( num , 1 , Integer: : sum ); ++ size3 ; } else { mid . merge ( num , 1 , Integer: : sum ); s += num ; } q . offer ( num ); if ( q . size () > m ) { int x = q . poll (); if ( lo . containsKey ( x )) { if ( lo . merge ( x , - 1 , Integer: : sum ) == 0 ) { lo . remove ( x ); } -- size1 ; } else if ( hi . containsKey ( x )) { if ( hi . merge ( x , - 1 , Integer: : sum ) == 0 ) { hi . remove ( x ); } -- size3 ; } else { if ( mid . merge ( x , - 1 , Integer: : sum ) == 0 ) { mid . remove ( x ); } s -= x ; } } for (; size1 > k ; -- size1 ) { int x = lo . lastKey (); if ( lo . merge ( x , - 1 , Integer: : sum ) == 0 ) { lo . remove ( x ); } mid . merge ( x , 1 , Integer: : sum ); s += x ; } for (; size3 > k ; -- size3 ) { int x = hi . firstKey (); if ( hi . merge ( x , - 1 , Integer: : sum ) == 0 ) { hi . remove ( x ); } mid . merge ( x , 1 , Integer: : sum ); s += x ; } for (; size1 < k && ! mid . isEmpty (); ++ size1 ) { int x = mid . firstKey (); if ( mid . merge ( x , - 1 , Integer: : sum ) == 0 ) { mid . remove ( x ); } s -= x ; lo . merge ( x , 1 , Integer: : sum ); } for (; size3 < k && ! mid . isEmpty (); ++ size3 ) { int x = mid . lastKey (); if ( mid . merge ( x , - 1 , Integer: : sum ) == 0 ) { mid . remove ( x ); } s -= x ; hi . merge ( x , 1 , Integer: : sum ); } } public int calculateMKAverage () { return q . size () < m ? - 1 : ( int ) ( s / ( q . size () - k * 2 )); } } /** * Your MKAverage object will be instantiated and called as such: * MKAverage obj = new MKAverage(m, k); * obj.addElement(num); * int param_2 = obj.calculateMKAverage(); */
```

### CPP

```cpp
class MKAverage { public: MKAverage ( int m , int k ) { this -> m = m ; this -> k = k ; } void addElement ( int num ) { if ( lo . empty () || num <= * lo . rbegin ()) { lo . insert ( num ); } else if ( hi . empty () || num >= * hi . begin ()) { hi . insert ( num ); } else { mid . insert ( num ); s += num ; } q . push ( num ); if ( q . size () > m ) { int x = q . front (); q . pop (); if ( lo . find ( x ) != lo . end ()) { lo . erase ( lo . find ( x )); } else if ( hi . find ( x ) != hi . end ()) { hi . erase ( hi . find ( x )); } else { mid . erase ( mid . find ( x )); s -= x ; } } while ( lo . size () > k ) { int x = * lo . rbegin (); lo . erase ( prev ( lo . end ())); mid . insert ( x ); s += x ; } while ( hi . size () > k ) { int x = * hi . begin (); hi . erase ( hi . begin ()); mid . insert ( x ); s += x ; } while ( lo . size () < k && mid . size ()) { int x = * mid . begin (); mid . erase ( mid . begin ()); s -= x ; lo . insert ( x ); } while ( hi . size () < k && mid . size ()) { int x = * mid . rbegin (); mid . erase ( prev ( mid . end ())); s -= x ; hi . insert ( x ); } } int calculateMKAverage () { return q . size () < m ? - 1 : s / ( q . size () - k * 2 ); } private: int m , k ; long long s = 0 ; queue < int > q ; multiset < int > lo , mid , hi ; }; /** * Your MKAverage object will be instantiated and called as such: * MKAverage* obj = new MKAverage(m, k); * obj->addElement(num); * int param_2 = obj->calculateMKAverage(); */
```

### Python

```python
from sortedcontainers import SortedList class MKAverage : def __init__ ( self , m : int , k : int ): self . m = m self . k = k self . s = 0 self . q = deque () self . lo = SortedList () self . mid = SortedList () self . hi = SortedList () def addElement ( self , num : int ) -> None : if not self . lo or num <= self . lo [ - 1 ]: self . lo . add ( num ) elif not self . hi or num >= self . hi [ 0 ]: self . hi . add ( num ) else : self . mid . add ( num ) self . s += num self . q . append ( num ) if len ( self . q ) > self . m : x = self . q . popleft () if x in self . lo : self . lo . remove ( x ) elif x in self . hi : self . hi . remove ( x ) else : self . mid . remove ( x ) self . s -= x while len ( self . lo ) > self . k : x = self . lo . pop () self . mid . add ( x ) self . s += x while len ( self . hi ) > self . k : x = self . hi . pop ( 0 ) self . mid . add ( x ) self . s += x while len ( self . lo ) < self . k and self . mid : x = self . mid . pop ( 0 ) self . lo . add ( x ) self . s -= x while len ( self . hi ) < self . k and self . mid : x = self . mid . pop () self . hi . add ( x ) self . s -= x def calculateMKAverage ( self ) -> int : return - 1 if len ( self . q ) < self . m else self . s // ( 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()
```