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 <= 105
1 <= k*2 < m
1 <= 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)$.
fromsortedcontainersimportSortedListclassMKAverage:def__init__(self,m:int,k:int):self.m=mself.k=kself.s=0self.q=deque()self.lo=SortedList()self.mid=SortedList()self.hi=SortedList()defaddElement(self,num:int)->None:ifnotself.loornum<=self.lo[-1]:self.lo.add(num)elifnotself.hiornum>=self.hi[0]:self.hi.add(num)else:self.mid.add(num)self.s+=numself.q.append(num)iflen(self.q)>self.m:x=self.q.popleft()ifxinself.lo:self.lo.remove(x)elifxinself.hi:self.hi.remove(x)else:self.mid.remove(x)self.s-=xwhilelen(self.lo)>self.k:x=self.lo.pop()self.mid.add(x)self.s+=xwhilelen(self.hi)>self.k:x=self.hi.pop(0)self.mid.add(x)self.s+=xwhilelen(self.lo)<self.kandself.mid:x=self.mid.pop(0)self.lo.add(x)self.s-=xwhilelen(self.hi)<self.kandself.mid:x=self.mid.pop()self.hi.add(x)self.s-=xdefcalculateMKAverage(self)->int:return-1iflen(self.q)<self.melseself.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()
classMKAverage{privateintm,k;privatelongs;privateintsize1,size3;privateDeque<Integer>q=newArrayDeque<>();privateTreeMap<Integer,Integer>lo=newTreeMap<>();privateTreeMap<Integer,Integer>mid=newTreeMap<>();privateTreeMap<Integer,Integer>hi=newTreeMap<>();publicMKAverage(intm,intk){this.m=m;this.k=k;}publicvoidaddElement(intnum){if(lo.isEmpty()||num<=lo.lastKey()){lo.merge(num,1,Integer::sum);++size1;}elseif(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){intx=q.poll();if(lo.containsKey(x)){if(lo.merge(x,-1,Integer::sum)==0){lo.remove(x);}--size1;}elseif(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){intx=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){intx=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){intx=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){intx=mid.lastKey();if(mid.merge(x,-1,Integer::sum)==0){mid.remove(x);}s-=x;hi.merge(x,1,Integer::sum);}}publicintcalculateMKAverage(){returnq.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(); */
classMKAverage{public:MKAverage(intm,intk){this->m=m;this->k=k;}voidaddElement(intnum){if(lo.empty()||num<=*lo.rbegin()){lo.insert(num);}elseif(hi.empty()||num>=*hi.begin()){hi.insert(num);}else{mid.insert(num);s+=num;}q.push(num);if(q.size()>m){intx=q.front();q.pop();if(lo.find(x)!=lo.end()){lo.erase(lo.find(x));}elseif(hi.find(x)!=hi.end()){hi.erase(hi.find(x));}else{mid.erase(mid.find(x));s-=x;}}while(lo.size()>k){intx=*lo.rbegin();lo.erase(prev(lo.end()));mid.insert(x);s+=x;}while(hi.size()>k){intx=*hi.begin();hi.erase(hi.begin());mid.insert(x);s+=x;}while(lo.size()<k&&mid.size()){intx=*mid.begin();mid.erase(mid.begin());s-=x;lo.insert(x);}while(hi.size()<k&&mid.size()){intx=*mid.rbegin();mid.erase(prev(mid.end()));s-=x;hi.insert(x);}}intcalculateMKAverage(){returnq.size()<m?-1:s/(q.size()-k*2);}private:intm,k;longlongs=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(); */
typeMKAveragestruct{lo,mid,hi*redblacktree.Treeq[]intm,k,sintsize1,size3int}funcConstructor(mint,kint)MKAverage{lo:=redblacktree.NewWithIntComparator()mid:=redblacktree.NewWithIntComparator()hi:=redblacktree.NewWithIntComparator()returnMKAverage{lo,mid,hi,[]int{},m,k,0,0,0}}func(this*MKAverage)AddElement(numint){merge:=func(rbt*redblacktree.Tree,key,valueint){ifv,ok:=rbt.Get(key);ok{nxt:=v.(int)+valueifnxt==0{rbt.Remove(key)}else{rbt.Put(key,nxt)}}else{rbt.Put(key,value)}}ifthis.lo.Empty()||num<=this.lo.Right().Key.(int){merge(this.lo,num,1)this.size1++}elseifthis.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)iflen(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--}elseif_,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-=xmerge(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-=xmerge(this.hi,x,1)}}func(this*MKAverage)CalculateMKAverage()int{iflen(this.q)<this.m{return-1}returnthis.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(); */
fromsortedcontainersimportSortedListclassMKAverage:def__init__(self,m:int,k:int):self.m=mself.k=kself.sl=SortedList()self.q=deque()self.s=0defaddElement(self,num:int)->None:self.q.append(num)iflen(self.q)==self.m:self.sl=SortedList(self.q)self.s=sum(self.sl[self.k:-self.k])eliflen(self.q)>self.m:i=self.sl.bisect_left(num)ifi<self.k:self.s+=self.sl[self.k-1]elifself.k<=i<=self.m-self.k:self.s+=numelse:self.s+=self.sl[self.m-self.k]self.sl.add(num)x=self.q.popleft()i=self.sl.bisect_left(x)ifi<self.k:self.s-=self.sl[self.k]elifself.k<=i<=self.m-self.k:self.s-=xelse:self.s-=self.sl[self.m-self.k]self.sl.remove(x)defcalculateMKAverage(self)->int:return-1iflen(self.sl)<self.melseself.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()