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:
- 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 <= 1051 <= k*2 < m1 <= num <= 105- At most
105 calls will be made to addElement and calculateMKAverage.
Solutions#
Solution 1: Ordered Set + Queue#
We can maintain the following data structures or variables:
- A queue $q$ of length $m$, where the head of the queue is the earliest added element, and the tail of the queue is the most recently added element;
- Three ordered sets, namely $lo$, $mid$, $hi$, where $lo$ and $hi$ store the smallest $k$ elements and the largest $k$ elements respectively, and $mid$ stores the remaining elements;
- A variable $s$, maintaining the sum of all elements in $mid$;
- Some programming languages (such as Java, Go) additionally maintain two variables $size1$ and $size3$, representing the number of elements in $lo$ and $hi$ respectively.
When calling the $addElement(num)$ function, perform the following operations in order:
- If $lo$ is empty, or $num \leq max(lo)$, then add $num$ to $lo$; otherwise if $hi$ is empty, or $num \geq min(hi)$, then add $num$ to $hi$; otherwise add $num$ to $mid$, and add the value of $num$ to $s$.
- Next, add $num$ to the queue $q$. If the length of the queue $q$ is greater than $m$ at this time, remove the head element $x$ from the queue $q$, then choose one of $lo$, $mid$ or $hi$ that contains $x$, and remove $x$ from this set. If the set is $mid$, subtract the value of $x$ from $s$.
- If the length of $lo$ is greater than $k$, then repeatedly remove the maximum value $max(lo)$ from $lo$, add $max(lo)$ to $mid$, and add the value of $max(lo)$ to $s$.
- If the length of $hi$ is greater than $k$, then repeatedly remove the minimum value $min(hi)$ from $hi$, add $min(hi)$ to $mid$, and add the value of $min(hi)$ to $s$.
- If the length of $lo$ is less than $k$ and $mid$ is not empty, then repeatedly remove the minimum value $min(mid)$ from $mid$, add $min(mid)$ to $lo$, and subtract the value of $min(mid)$ from $s$.
- If the length of $hi$ is less than $k$ and $mid$ is not empty, then repeatedly remove the maximum value $max(mid)$ from $mid$, add $max(mid)$ to $hi$, and subtract the value of $max(mid)$ from $s$.
When calling the $calculateMKAverage()$ function, if the length of $q$ is less than $m$, return $-1$, otherwise return $\frac{s}{m - 2k}$.
In terms of time complexity, each call to the $addElement(num)$ function has a time complexity of $O(\log m)$, and each call to the $calculateMKAverage()$ function has a time complexity of $O(1)$. The space complexity is $O(m)$.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
| from sortedcontainers import SortedList
class MKAverage:
def __init__(self, m: int, k: int):
self.m = m
self.k = k
self.s = 0
self.q = deque()
self.lo = SortedList()
self.mid = SortedList()
self.hi = SortedList()
def addElement(self, num: int) -> None:
if not self.lo or num <= self.lo[-1]:
self.lo.add(num)
elif not self.hi or num >= self.hi[0]:
self.hi.add(num)
else:
self.mid.add(num)
self.s += num
self.q.append(num)
if len(self.q) > self.m:
x = self.q.popleft()
if x in self.lo:
self.lo.remove(x)
elif x in self.hi:
self.hi.remove(x)
else:
self.mid.remove(x)
self.s -= x
while len(self.lo) > self.k:
x = self.lo.pop()
self.mid.add(x)
self.s += x
while len(self.hi) > self.k:
x = self.hi.pop(0)
self.mid.add(x)
self.s += x
while len(self.lo) < self.k and self.mid:
x = self.mid.pop(0)
self.lo.add(x)
self.s -= x
while len(self.hi) < self.k and self.mid:
x = self.mid.pop()
self.hi.add(x)
self.s -= x
def calculateMKAverage(self) -> int:
return -1 if len(self.q) < self.m else self.s // (self.m - 2 * self.k)
# Your MKAverage object will be instantiated and called as such:
# obj = MKAverage(m, k)
# obj.addElement(num)
# param_2 = obj.calculateMKAverage()
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
| class MKAverage {
private int m, k;
private long s;
private int size1, size3;
private Deque<Integer> q = new ArrayDeque<>();
private TreeMap<Integer, Integer> lo = new TreeMap<>();
private TreeMap<Integer, Integer> mid = new TreeMap<>();
private TreeMap<Integer, Integer> hi = new TreeMap<>();
public MKAverage(int m, int k) {
this.m = m;
this.k = k;
}
public void addElement(int num) {
if (lo.isEmpty() || num <= lo.lastKey()) {
lo.merge(num, 1, Integer::sum);
++size1;
} else if (hi.isEmpty() || num >= hi.firstKey()) {
hi.merge(num, 1, Integer::sum);
++size3;
} else {
mid.merge(num, 1, Integer::sum);
s += num;
}
q.offer(num);
if (q.size() > m) {
int x = q.poll();
if (lo.containsKey(x)) {
if (lo.merge(x, -1, Integer::sum) == 0) {
lo.remove(x);
}
--size1;
} else if (hi.containsKey(x)) {
if (hi.merge(x, -1, Integer::sum) == 0) {
hi.remove(x);
}
--size3;
} else {
if (mid.merge(x, -1, Integer::sum) == 0) {
mid.remove(x);
}
s -= x;
}
}
for (; size1 > k; --size1) {
int x = lo.lastKey();
if (lo.merge(x, -1, Integer::sum) == 0) {
lo.remove(x);
}
mid.merge(x, 1, Integer::sum);
s += x;
}
for (; size3 > k; --size3) {
int x = hi.firstKey();
if (hi.merge(x, -1, Integer::sum) == 0) {
hi.remove(x);
}
mid.merge(x, 1, Integer::sum);
s += x;
}
for (; size1 < k && !mid.isEmpty(); ++size1) {
int x = mid.firstKey();
if (mid.merge(x, -1, Integer::sum) == 0) {
mid.remove(x);
}
s -= x;
lo.merge(x, 1, Integer::sum);
}
for (; size3 < k && !mid.isEmpty(); ++size3) {
int x = mid.lastKey();
if (mid.merge(x, -1, Integer::sum) == 0) {
mid.remove(x);
}
s -= x;
hi.merge(x, 1, Integer::sum);
}
}
public int calculateMKAverage() {
return q.size() < m ? -1 : (int) (s / (q.size() - k * 2));
}
}
/**
* Your MKAverage object will be instantiated and called as such:
* MKAverage obj = new MKAverage(m, k);
* obj.addElement(num);
* int param_2 = obj.calculateMKAverage();
*/
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
| class MKAverage {
public:
MKAverage(int m, int k) {
this->m = m;
this->k = k;
}
void addElement(int num) {
if (lo.empty() || num <= *lo.rbegin()) {
lo.insert(num);
} else if (hi.empty() || num >= *hi.begin()) {
hi.insert(num);
} else {
mid.insert(num);
s += num;
}
q.push(num);
if (q.size() > m) {
int x = q.front();
q.pop();
if (lo.find(x) != lo.end()) {
lo.erase(lo.find(x));
} else if (hi.find(x) != hi.end()) {
hi.erase(hi.find(x));
} else {
mid.erase(mid.find(x));
s -= x;
}
}
while (lo.size() > k) {
int x = *lo.rbegin();
lo.erase(prev(lo.end()));
mid.insert(x);
s += x;
}
while (hi.size() > k) {
int x = *hi.begin();
hi.erase(hi.begin());
mid.insert(x);
s += x;
}
while (lo.size() < k && mid.size()) {
int x = *mid.begin();
mid.erase(mid.begin());
s -= x;
lo.insert(x);
}
while (hi.size() < k && mid.size()) {
int x = *mid.rbegin();
mid.erase(prev(mid.end()));
s -= x;
hi.insert(x);
}
}
int calculateMKAverage() {
return q.size() < m ? -1 : s / (q.size() - k * 2);
}
private:
int m, k;
long long s = 0;
queue<int> q;
multiset<int> lo, mid, hi;
};
/**
* Your MKAverage object will be instantiated and called as such:
* MKAverage* obj = new MKAverage(m, k);
* obj->addElement(num);
* int param_2 = obj->calculateMKAverage();
*/
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
| type MKAverage struct {
lo, mid, hi *redblacktree.Tree
q []int
m, k, s int
size1, size3 int
}
func Constructor(m int, k int) MKAverage {
lo := redblacktree.NewWithIntComparator()
mid := redblacktree.NewWithIntComparator()
hi := redblacktree.NewWithIntComparator()
return MKAverage{lo, mid, hi, []int{}, m, k, 0, 0, 0}
}
func (this *MKAverage) AddElement(num int) {
merge := func(rbt *redblacktree.Tree, key, value int) {
if v, ok := rbt.Get(key); ok {
nxt := v.(int) + value
if nxt == 0 {
rbt.Remove(key)
} else {
rbt.Put(key, nxt)
}
} else {
rbt.Put(key, value)
}
}
if this.lo.Empty() || num <= this.lo.Right().Key.(int) {
merge(this.lo, num, 1)
this.size1++
} else if this.hi.Empty() || num >= this.hi.Left().Key.(int) {
merge(this.hi, num, 1)
this.size3++
} else {
merge(this.mid, num, 1)
this.s += num
}
this.q = append(this.q, num)
if len(this.q) > this.m {
x := this.q[0]
this.q = this.q[1:]
if _, ok := this.lo.Get(x); ok {
merge(this.lo, x, -1)
this.size1--
} else if _, ok := this.hi.Get(x); ok {
merge(this.hi, x, -1)
this.size3--
} else {
merge(this.mid, x, -1)
this.s -= x
}
}
for ; this.size1 > this.k; this.size1-- {
x := this.lo.Right().Key.(int)
merge(this.lo, x, -1)
merge(this.mid, x, 1)
this.s += x
}
for ; this.size3 > this.k; this.size3-- {
x := this.hi.Left().Key.(int)
merge(this.hi, x, -1)
merge(this.mid, x, 1)
this.s += x
}
for ; this.size1 < this.k && !this.mid.Empty(); this.size1++ {
x := this.mid.Left().Key.(int)
merge(this.mid, x, -1)
this.s -= x
merge(this.lo, x, 1)
}
for ; this.size3 < this.k && !this.mid.Empty(); this.size3++ {
x := this.mid.Right().Key.(int)
merge(this.mid, x, -1)
this.s -= x
merge(this.hi, x, 1)
}
}
func (this *MKAverage) CalculateMKAverage() int {
if len(this.q) < this.m {
return -1
}
return this.s / (this.m - 2*this.k)
}
/**
* Your MKAverage object will be instantiated and called as such:
* obj := Constructor(m, k);
* obj.AddElement(num);
* param_2 := obj.CalculateMKAverage();
*/
|
Solution 2#
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
| from sortedcontainers import SortedList
class MKAverage:
def __init__(self, m: int, k: int):
self.m = m
self.k = k
self.sl = SortedList()
self.q = deque()
self.s = 0
def addElement(self, num: int) -> None:
self.q.append(num)
if len(self.q) == self.m:
self.sl = SortedList(self.q)
self.s = sum(self.sl[self.k : -self.k])
elif len(self.q) > self.m:
i = self.sl.bisect_left(num)
if i < self.k:
self.s += self.sl[self.k - 1]
elif self.k <= i <= self.m - self.k:
self.s += num
else:
self.s += self.sl[self.m - self.k]
self.sl.add(num)
x = self.q.popleft()
i = self.sl.bisect_left(x)
if i < self.k:
self.s -= self.sl[self.k]
elif self.k <= i <= self.m - self.k:
self.s -= x
else:
self.s -= self.sl[self.m - self.k]
self.sl.remove(x)
def calculateMKAverage(self) -> int:
return -1 if len(self.sl) < self.m else self.s // (self.m - self.k * 2)
# Your MKAverage object will be instantiated and called as such:
# obj = MKAverage(m, k)
# obj.addElement(num)
# param_2 = obj.calculateMKAverage()
|
