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:
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.
Remove the smallest k elements and the largest k elements from the container.
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
class MKAverage {
int m;
int k;
queue<int> window;
multiset<int> smaller_nums;
multiset<int> larger_nums;
multiset<int> middles;
long middles_sum = 0;
void InitializeMultiset() {
vector<int> nums;
queue<int> current_window = window;
while (!current_window.empty()) {
nums.push_back(current_window.front());
current_window.pop();
}
sort(nums.begin(), nums.end());
for (int i = 0; i < m; i++) {
if (i < k) {
smaller_nums.insert(nums[i]);
} else if (i > m - k - 1) {
larger_nums.insert(nums[i]);
} else {
middles_sum += nums[i];
middles.insert(nums[i]);
}
}
}
public:
MKAverage(int m, int k) {
this->m = m;
this->k = k;
}
void AddElement(int val) {
window.push(val);
if (middles.empty()) {
// It only run once.
if (window.size() == m) {
InitializeMultiset();
}
} else if (val < *smaller_nums.rbegin()) {
int middle_candidate = *smaller_nums.rbegin();
smaller_nums.erase(--smaller_nums.end());
smaller_nums.insert(val);
middles.insert(middle_candidate);
middles_sum += middle_candidate;
} else if (val > *larger_nums.begin()) {
int middle_candidate = *larger_nums.begin();
larger_nums.erase(larger_nums.begin());
larger_nums.insert(val);
middles.insert(middle_candidate);
middles_sum += middle_candidate;
} else {
middles.insert(val);
middles_sum += val;
}
}
void RemoveElement() {
int val = window.front();
window.pop();
if (val <= *smaller_nums.rbegin()) {
smaller_nums.erase(smaller_nums.lower_bound(val));
int small_candidate = *middles.begin();
smaller_nums.insert(small_candidate);
middles.erase(middles.begin());
middles_sum -= small_candidate;
} else if (val >= * larger_nums.begin()) {
larger_nums.erase(larger_nums.lower_bound(val));
int large_candidate = *middles.rbegin();
larger_nums.insert(large_candidate);
middles.erase(--middles.end());
middles_sum -= large_candidate;
} else {
middles.erase(middles.lower_bound(val));
middles_sum -= val;
}
}
void addElement(int num) {
AddElement(num);
if (window.size() > m) {
RemoveElement();
}
}
int calculateMKAverage() {
return window.size() < m ? -1 : middles_sum / middles.size();
}
};