# 1825. Finding MK Average

## Description

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
[[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.

## Solutions

Solution 1: Ordered Set + Queue

We can maintain the following data structures or variables:

• A queue $q$ of length $m$, where the head of the queue is the earliest added element, and the tail of the queue is the most recently added element;
• Three ordered sets, namely $lo$, $mid$, $hi$, where $lo$ and $hi$ store the smallest $k$ elements and the largest $k$ elements respectively, and $mid$ stores the remaining elements;
• A variable $s$, maintaining the sum of all elements in $mid$;
• Some programming languages (such as Java, Go) additionally maintain two variables $size1$ and $size3$, representing the number of elements in $lo$ and $hi$ respectively.

When calling the $addElement(num)$ function, perform the following operations in order:

1. If $lo$ is empty, or $num \leq max(lo)$, then add $num$ to $lo$; otherwise if $hi$ is empty, or $num \geq min(hi)$, then add $num$ to $hi$; otherwise add $num$ to $mid$, and add the value of $num$ to $s$.
2. Next, add $num$ to the queue $q$. If the length of the queue $q$ is greater than $m$ at this time, remove the head element $x$ from the queue $q$, then choose one of $lo$, $mid$ or $hi$ that contains $x$, and remove $x$ from this set. If the set is $mid$, subtract the value of $x$ from $s$.
3. If the length of $lo$ is greater than $k$, then repeatedly remove the maximum value $max(lo)$ from $lo$, add $max(lo)$ to $mid$, and add the value of $max(lo)$ to $s$.
4. If the length of $hi$ is greater than $k$, then repeatedly remove the minimum value $min(hi)$ from $hi$, add $min(hi)$ to $mid$, and add the value of $min(hi)$ to $s$.
5. If the length of $lo$ is less than $k$ and $mid$ is not empty, then repeatedly remove the minimum value $min(mid)$ from $mid$, add $min(mid)$ to $lo$, and subtract the value of $min(mid)$ from $s$.
6. If the length of $hi$ is less than $k$ and $mid$ is not empty, then repeatedly remove the maximum value $max(mid)$ from $mid$, add $max(mid)$ to $hi$, and subtract the value of $max(mid)$ from $s$.

When calling the $calculateMKAverage()$ function, if the length of $q$ is less than $m$, return $-1$, otherwise return $\frac{s}{m - 2k}$.

In terms of time complexity, each call to the $addElement(num)$ function has a time complexity of $O(\log m)$, and each call to the $calculateMKAverage()$ function has a time complexity of $O(1)$. The space complexity is $O(m)$.

• 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;
}

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);
* int param_2 = obj.calculateMKAverage();
*/

• class MKAverage {
public:
MKAverage(int m, int k) {
this->m = m;
this->k = k;
}

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);
* int param_2 = obj->calculateMKAverage();
*/

• 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]:
elif not self.hi or num >= self.hi[0]:
else:
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.s += x
while len(self.hi) > self.k:
x = self.hi.pop(0)
self.s += x
while len(self.lo) < self.k and self.mid:
x = self.mid.pop(0)
self.s -= x
while len(self.hi) < self.k and self.mid:
x = self.mid.pop()
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)
# param_2 = obj.calculateMKAverage()


• type MKAverage struct {
lo, mid, hi  *redblacktree.Tree
q            []int
m, k, s      int
size1, size3 int
}

func Constructor(m int, k int) MKAverage {
lo := redblacktree.NewWithIntComparator()
mid := redblacktree.NewWithIntComparator()
hi := redblacktree.NewWithIntComparator()
return MKAverage{lo, mid, hi, []int{}, m, k, 0, 0, 0}
}

func (this *MKAverage) AddElement(num int) {
merge := func(rbt *redblacktree.Tree, key, value int) {
if v, ok := rbt.Get(key); ok {
nxt := v.(int) + value
if nxt == 0 {
rbt.Remove(key)
} else {
rbt.Put(key, nxt)
}
} else {
rbt.Put(key, value)
}
}

if this.lo.Empty() || num <= this.lo.Right().Key.(int) {
merge(this.lo, num, 1)
this.size1++
} else if this.hi.Empty() || num >= this.hi.Left().Key.(int) {
merge(this.hi, num, 1)
this.size3++
} else {
merge(this.mid, num, 1)
this.s += num
}
this.q = append(this.q, num)
if len(this.q) > this.m {
x := this.q[0]
this.q = this.q[1:]
if _, ok := this.lo.Get(x); ok {
merge(this.lo, x, -1)
this.size1--
} else if _, ok := this.hi.Get(x); ok {
merge(this.hi, x, -1)
this.size3--
} else {
merge(this.mid, x, -1)
this.s -= x
}
}
for ; this.size1 > this.k; this.size1-- {
x := this.lo.Right().Key.(int)
merge(this.lo, x, -1)
merge(this.mid, x, 1)
this.s += x
}
for ; this.size3 > this.k; this.size3-- {
x := this.hi.Left().Key.(int)
merge(this.hi, x, -1)
merge(this.mid, x, 1)
this.s += x
}
for ; this.size1 < this.k && !this.mid.Empty(); this.size1++ {
x := this.mid.Left().Key.(int)
merge(this.mid, x, -1)
this.s -= x
merge(this.lo, x, 1)
}
for ; this.size3 < this.k && !this.mid.Empty(); this.size3++ {
x := this.mid.Right().Key.(int)
merge(this.mid, x, -1)
this.s -= x
merge(this.hi, x, 1)
}
}

func (this *MKAverage) CalculateMKAverage() int {
if len(this.q) < this.m {
return -1
}
return this.s / (this.m - 2*this.k)
}

/**
* Your MKAverage object will be instantiated and called as such:
* obj := Constructor(m, k);