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 1: Multiset * 3

Use three multiset to track the left part (smallest k elements), right part (largest k elements) and mid (middle part of m – 2*k elements).

Time complexity: addElememt: O(logn), average: O(1)
Space complexity: O(n)

class MKAverage {
public:
  MKAverage(int m, int k): 
    sum(0), m(m), k(k), n(m - 2*k) {}

  void addElement(int num) {
    if (q.size() == m) {      
      remove(q.front());
      q.pop();
    }
    q.push(num);
    add(num);
  }

  int calculateMKAverage() {    
    return (q.size() < m) ? -1 : sum / n;
  }
private:
  void add(int x) {
    left.insert(x);
    
    if (left.size() > k) {
      auto it = prev(end(left));
      sum += *it;
      mid.insert(*it);      
      left.erase(it);
    }
    
    if (mid.size() > n) {
      auto it = prev(end(mid));
      sum -= *it; 
      right.insert(*it);
      mid.erase(it);
    }
  }
  
  void remove(int x) {
    if (x <= *rbegin(left)) {
      left.erase(left.find(x));
    } else if (x <= *rbegin(mid)) {
      sum -= x;
      mid.erase(mid.find(x));
    } else {
      right.erase(right.find(x));
    }
    
    if (left.size() < k) {
      auto it = begin(mid);
      sum -= *it;
      left.insert(*it);
      mid.erase(it);
    }
    
    if (mid.size() < n) {
      auto it = begin(right);
      sum += *it;
      mid.insert(*it);
      right.erase(it);
    }
  }
  
  queue<int> q;
  multiset<int> left, mid, right;  
  long sum;
  const int m;
  const int k;
  const int n;
};

The post 花花酱 LeetCode 1825. Finding MK Average appeared first on Huahua's Tech Road.