##### Welcome to Subscribe On Youtube

Formatted question description: https://leetcode.ca/all/1825.html

Hard

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

## Solution

Use a queue to store the elements, and use three tree maps left, mid and right to store the smallest k elements, the middle elements and the largest k elements, respectively.

For the constructor, initialize m, k, the queue and the three tree maps.

For addElement, first check whether the queue’s size is already m. If this is the case, poll one element from the queue, and update the three tree maps accordingly. Then offer the new element to the queue, and use the new element to update the three tree maps accordingly. At the same time maintain the sum of the elements in mid.

For calculateMKAverage, first check the queue’s size, and then calculate the MKAverage accordingly.

• class MKAverage {
int m;
int k;
Queue<Integer> queue;
TreeMap<Integer, Integer> left;
TreeMap<Integer, Integer> mid;
TreeMap<Integer, Integer> right;
int leftSize, midSize, rightSize;
long sum;

public MKAverage(int m, int k) {
this.m = m;
this.k = k;
left = new TreeMap<Integer, Integer>();
mid = new TreeMap<Integer, Integer>();
right = new TreeMap<Integer, Integer>();
}

if (queue.size() == m) {
int remove = queue.poll();
if (left.containsKey(remove)) {
left.put(remove, left.get(remove) - 1);
if (left.get(remove) == 0)
left.remove(remove);
leftSize--;
} else if (mid.containsKey(remove)) {
mid.put(remove, mid.get(remove) - 1);
if (mid.get(remove) == 0)
mid.remove(remove);
midSize--;
sum -= remove;
} else {
right.put(remove, right.get(remove) - 1);
if (right.get(remove) == 0)
right.remove(remove);
rightSize--;
}
while (midSize > 0 && leftSize < k) {
int first = mid.firstKey();
left.put(first, left.getOrDefault(first, 0) + 1);
mid.put(first, mid.get(first) - 1);
if (mid.get(first) == 0)
mid.remove(first);
leftSize++;
midSize--;
sum -= first;
}
while (rightSize > 0 && midSize < m - 2 * k) {
int first = right.firstKey();
mid.put(first, mid.getOrDefault(first, 0) + 1);
right.put(first, right.get(first) - 1);
if (right.get(first) == 0)
right.remove(first);
midSize++;
rightSize--;
sum += first;
}
}
queue.offer(num);
if (leftSize < k || num <= left.lastKey()) {
left.put(num, left.getOrDefault(num, 0) + 1);
leftSize++;
} else if (midSize < m - 2 * k || num <= mid.lastKey()) {
mid.put(num, mid.getOrDefault(num, 0) + 1);
midSize++;
sum += num;
} else {
right.put(num, right.getOrDefault(num, 0) + 1);
rightSize++;
}
while (leftSize > k) {
int last = left.lastKey();
left.put(last, left.get(last) - 1);
if (left.get(last) == 0)
left.remove(last);
mid.put(last, mid.getOrDefault(last, 0) + 1);
leftSize--;
midSize++;
sum += last;
}
while (midSize > m - 2 * k) {
int last = mid.lastKey();
mid.put(last, mid.get(last) - 1);
if (mid.get(last) == 0)
mid.remove(last);
right.put(last, right.getOrDefault(last, 0) + 1);
midSize--;
rightSize++;
sum -= last;
}
}

public int calculateMKAverage() {
return queue.size() < m ? -1 : (int) (sum / (m - 2 * k));
}
}

/**
* Your MKAverage object will be instantiated and called as such:
* MKAverage obj = new MKAverage(m, k);
* int param_2 = obj.calculateMKAverage();
*/

############

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();
*/

• 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()

############

# 1825. Finding MK Average
# https://leetcode.com/problems/finding-mk-average

class MKAverage:

def __init__(self, m: int, k: int):
self.m = m
self.k = k
self.A = []

def addElement(self, num: int) -> None:
self.A.append(num)

def calculateMKAverage(self) -> int:
if len(self.A) < self.m: return -1

m, k = self.m, self.k
A = self.A[:]

if len(A) > m:
A = self.A[len(A) - m:]
A.sort()

total = 0
for i in range(k, len(A) - k):
total += A[i]

# Your MKAverage object will be instantiated and called as such:
# obj = MKAverage(m, k)
# 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();
*/

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