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 <= 105

• 1 <= k*2 < m

• 1 <= num <= 105

• At most 105 calls will be made to addElement and calculateMKAverage.

Solution

The idea is to maintain 3 different map to records the value and their count. Top, Middle, Button

Whenever an element is added, first add it into Middle. And then, we start to adjust the 3 map according to satified: 1. Top has the largest k element 2. Bottom has the smallest k element

There are 2 different cases for (midle <-> top) or (middle <-> bottom)

• Top or buttom size is less than k,

• Swap is needed for(top<->middle): Top's firstKey is smaller than middle.lastKey

• Swap is needed for (middle<->button): button's lastKey is greater than middle.firstKey

If it is not satisfied, we will need to migrate the last/first element of Top/Buttom map, to the Middle Map.

We can take advantage of TreeMap. TreeMap provides API to access firstKey or lastKey.

    TreeMap<Integer, Integer> bottom = new TreeMap<>();
    TreeMap<Integer, Integer> middle = new TreeMap<>();
    TreeMap<Integer, Integer> top = new TreeMap<>();
    int m = 0;
    int k = 0;
    int middleSum = 0;
    int topCount = 0;
    int bottomCount = 0;
    Queue<Integer> queue = new LinkedList<>();
    
    public MKAverage(int m, int k) {
        this.m = m;
        this.k = k;
    }
    
    public void addElement(int num) {
        if(queue.size() == m) {
            int rem = queue.poll();
            if(top.containsKey(rem)){
                remove(top, rem);
                topCount--;
            }else if(middle.containsKey(rem)){
                remove(middle, rem);
                middleSum -= rem;
            }else{
                bottomCount--;
                remove(bottom, rem);
            }
        }
        queue.add(num);
        middle.put(num, middle.getOrDefault(num, 0) + 1);
        middleSum += num;
        //adjust
        // top is not enough
        while(topCount < k && !middle.isEmpty()){
            topCount++;
            Integer middleLast = middle.lastKey();
            remove(middle, middleLast);
            middleSum -= middleLast;
            top.put(middleLast, top.getOrDefault(middleLast, 0) + 1);
        }
        //top.first is smaller than middle.last, swap is requried
        while(!top.isEmpty() && !middle.isEmpty() && middle.lastKey() > top.firstKey()){
            Integer middleLast = middle.lastKey();
            Integer topFirst = top.firstKey();
            remove(middle, middleLast);
            remove(top, topFirst);
            middleSum -= middleLast;
            middleSum += topFirst;
            middle.put(topFirst, middle.getOrDefault(topFirst, 0) + 1);          
            top.put(middleLast, top.getOrDefault(middleLast, 0) + 1);
        }
        //bottom is not enough
        while(bottomCount < k && !middle.isEmpty()){
            bottomCount++;
            Integer middleFirst = middle.firstKey();
            remove(middle, middleFirst);
            middleSum -= middleFirst;
            bottom.put(middleFirst, bottom.getOrDefault(middleFirst, 0) + 1);
        }
        //middle.first is smaller than bottom.last, swap is requried
        while(!bottom.isEmpty() && !middle.isEmpty() && middle.firstKey() < bottom.lastKey()){
            Integer bottomLast = bottom.lastKey();
            Integer middleFirst = middle.firstKey();
            remove(middle, middleFirst);
            remove(bottom, bottomLast);
            middleSum -= middleFirst;
            middleSum += bottomLast;
            middle.put(bottomLast, middle.getOrDefault(bottomLast, 0) + 1);          
            bottom.put(middleFirst, bottom.getOrDefault(middleFirst, 0) + 1);
        }
        
    }
    
    public void remove(TreeMap<Integer, Integer> map, int num) {
        map.put(num, map.get(num)-1); 
        if(map.get(num) == 0){
                map.remove(num);
        }
    }
    
    
    public int calculateMKAverage() {
        if(queue.size() != m) return -1;
        
        return (int)(middleSum / (m - 2 * k));
    }
}