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
myou should consider the MKAverage to be-1. Otherwise, copy the lastmelements of the stream to a separate container. - Remove the smallest
kelements and the largestkelements 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 integersmandk.void addElement(int num)Inserts a new elementnuminto 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 <= 1051 <= k*2 < m1 <= num <= 105- At most
105calls will be made toaddElementandcalculateMKAverage.
Companies: Google
Related Topics:
Design, Queue, Heap (Priority Queue), Data Stream, Ordered Set
Similar Questions:
- Find Median from Data Stream (Hard)
- Kth Largest Element in a Stream (Easy)
- Sequentially Ordinal Rank Tracker (Hard)
Hints:
- At each query, try to save and update the sum of the elements needed to calculate MKAverage.
- You can use BSTs for fast insertion and deletion of the elements.
- Use 3
multiset<int>to track the topk, bottomkand middle elements. - Use
queue<int> qto track the currentmnumbers. - Use
sumto track the sum of numbers inmid.
// OJ: https://leetcode.com/problems/finding-mk-average/
// Author: github.com/lzl124631x
// Time:
// MKAverage: O(1)
// addElement: O(logM)
// calculateMKAverage: O(1)
// Space: O(M)
class MKAverage {
multiset<int> top, bot, mid;
queue<int> q;
long sum = 0, m, k;
public:
MKAverage(int m, int k) : m(m), k(k) {}
void addElement(int n) {
if (q.size() < m) mid.insert(n); // when there are less than `m` numbers, always insert into `mid`.
q.push(n);
if (q.size() == m) {
// when we reached exactly `m` numbers, we nudge numbers from `mid` to `top` and `bot`, and calculate `sum`.
for (int i = 0; i < k; ++i) {
bot.insert(*mid.begin());
mid.erase(mid.begin());
}
for (int i = 0; i < k; ++i) {
top.insert(*mid.rbegin());
mid.erase(prev(mid.end()));
}
for (int x : mid) sum += x;
} else if (q.size() > m) {
// when we've seen more than `m` numbers. We first add the new number `n` to where it should belong.
// If we add `n` to `top` or `bot`, we balance them with `mid` to make sure `top` and `bot` have exactly `k` numbers
if (n < *bot.rbegin()) {
bot.insert(n);
int x = *bot.rbegin();
bot.erase(prev(bot.end()));
mid.insert(x);
sum += x;
} else if (n > *top.begin()) {
top.insert(n);
int x = *top.begin();
top.erase(top.begin());
mid.insert(x);
sum += x;
} else {
mid.insert(n);
sum += n;
}
// Then we remove the number `rm` that is no longer one of the latest `m` numbers.
int rm = q.front();
q.pop();
auto it = mid.find(rm);
if (it != mid.end()) { // first try removing from `mid`, then `top` or `bot`.
mid.erase(it);
sum -= rm;
} else {
it = top.find(rm);
if (it != top.end()) {
top.erase(it);
} else {
bot.erase(bot.find(rm));
}
}
// Lastly, balance `top` and `bot` if needed
if (bot.size() < k) {
int x = *mid.begin();
bot.insert(x);
mid.erase(mid.begin());
sum -= x;
} else if (top.size() < k) {
int x = *mid.rbegin();
top.insert(x);
mid.erase(prev(mid.end()));
sum -= x;
}
}
}
int calculateMKAverage() {
return q.size() == m ? (sum / (m - 2 * k)) : -1;
}
};Or
// OJ: https://leetcode.com/problems/finding-mk-average
// Author: github.com/lzl124631x
// Time:
// MKAverage: O(1)
// addElement: O(logM)
// calculateMKAverage: O(1)
// Space: O(M)
class MKAverage {
multiset<int> top, mid, bot;
queue<int> q;
long m, k, sum = 0;
public:
MKAverage(int m, int k) : m(m), k(k) {
}
void addElement(int n) {
q.push(n);
if (q.size() > m) {
int rm = q.front();
q.pop();
if (*top.begin() <= rm) {
top.erase(top.find(rm));
rm = *mid.rbegin();
top.insert(rm);
}
if (*mid.begin() <= rm) {
mid.erase(mid.find(rm));
sum -= rm;
rm = *bot.rbegin();
mid.insert(rm);
sum += rm;
}
bot.erase(bot.find(rm));
}
if (top.size() < k || n > *top.begin()) {
top.insert(n);
if (top.size() <= k) return;
n = *top.begin();
top.erase(top.begin());
}
if (mid.size() < m - 2 * k || n > *mid.begin()) {
mid.insert(n);
sum += n;
if (mid.size() <= m - 2 * k) return;
n = *mid.begin();
sum -= n;
mid.erase(mid.begin());
}
bot.insert(n);
}
int calculateMKAverage() {
if (q.size() < m) return -1;
return sum / (m - 2 * k);
}
};Try BIT: https://leetcode.com/problems/finding-mk-average/discuss/1152438/Python3-Fenwick-tree