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

Level

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

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;
        queue = new LinkedList<Integer>();
        left = new TreeMap<Integer, Integer>();
        mid = new TreeMap<Integer, Integer>();
        right = new TreeMap<Integer, Integer>();
    }

    public void addElement(int num) {
        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);
 * obj.addElement(num);
 * int param_2 = obj.calculateMKAverage();
 */

All Problems

All Solutions