---
description: MultiSet
---

# LeetCode 1825. Finding MK Average

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 <= 10^5`
* `1 <= k*2 < m`
* `1 <= num <= 10^5`
* At most `10^5` calls will be made to `addElement` and `calculateMKAverage`.

## Solution

[English Version in Youtube](https://youtu.be/tHMWcAikcnY)

[中文版解答Youtube Link](https://youtu.be/Nr_XX-joVm4)

[中文版解答Bilibili Link](https://www.bilibili.com/video/BV1uy4y147Pi/)

```
class MKAverage {
    int m;
    int k;
    queue<int> window;
    
    multiset<int> smaller_nums;
    multiset<int> larger_nums;
    multiset<int> middles;
    
    long middles_sum = 0;
    
    void InitializeMultiset() {
        vector<int> nums;
        queue<int> current_window = window;
        while (!current_window.empty()) {
            nums.push_back(current_window.front());
            current_window.pop();
        }
        sort(nums.begin(), nums.end());
        for (int i = 0; i < m; i++) {
            if (i < k) {
                smaller_nums.insert(nums[i]);
            } else if (i > m - k - 1) {
                larger_nums.insert(nums[i]);
            } else {
                middles_sum += nums[i];
                middles.insert(nums[i]);
            }
        }
    }
public:
    MKAverage(int m, int k) {
        this->m = m;
        this->k = k;
    }

    void AddElement(int val) {
        window.push(val);
        if (middles.empty()) {
            // It only run once.
            if (window.size() == m) {
                InitializeMultiset();
            }
        } else if (val < *smaller_nums.rbegin()) {
            int middle_candidate = *smaller_nums.rbegin();
            smaller_nums.erase(--smaller_nums.end());
            smaller_nums.insert(val);
            middles.insert(middle_candidate);
            middles_sum += middle_candidate;
        } else if (val > *larger_nums.begin()) {
            int middle_candidate = *larger_nums.begin();
            larger_nums.erase(larger_nums.begin());
            larger_nums.insert(val);
            middles.insert(middle_candidate);
            middles_sum += middle_candidate;            
        } else {
            middles.insert(val);
            middles_sum += val;
        }
    }

    void RemoveElement() {
        int val = window.front();
        window.pop();
        if (val <= *smaller_nums.rbegin()) {
            smaller_nums.erase(smaller_nums.lower_bound(val));
            int small_candidate = *middles.begin();
            smaller_nums.insert(small_candidate);
            middles.erase(middles.begin());
            middles_sum -= small_candidate;
        } else if (val >= * larger_nums.begin()) {
            larger_nums.erase(larger_nums.lower_bound(val));
            int large_candidate = *middles.rbegin();
            larger_nums.insert(large_candidate);
            middles.erase(--middles.end());
            middles_sum -= large_candidate;
        } else {
            middles.erase(middles.lower_bound(val));
            middles_sum -= val;
        }
    }

    void addElement(int num) {
        AddElement(num);
        if (window.size() > m) {
            RemoveElement();
        }
    }

    int calculateMKAverage() {
        return window.size() < m ? -1 : middles_sum / middles.size();
    }
};
```
