You have an infinite number of stacks arranged in a row and numbered (left to right) from 0, each of the stacks has the same maximum capacity.
Implement the DinnerPlates class:
DinnerPlates(int capacity) Initializes the object with the maximum capacity of the stacks capacity.
void push(int val) Pushes the given integer val into the leftmost stack with a size less than capacity.
int pop() Returns the value at the top of the rightmost non-empty stack and removes it from that stack, and returns -1 if all the stacks are empty.
int popAtStack(int index) Returns the value at the top of the stack with the given index index and removes it from that stack or returns -1 if the stack with that given index is empty.
Example 1:
Input
["DinnerPlates", "push", "push", "push", "push", "push", "popAtStack", "push", "push", "popAtStack", "popAtStack", "pop", "pop", "pop", "pop", "pop"]
[[2], [1], [2], [3], [4], [5], [0], [20], [21], [0], [2], [], [], [], [], []]
Output
[null, null, null, null, null, null, 2, null, null, 20, 21, 5, 4, 3, 1, -1]
Explanation:
DinnerPlates D = DinnerPlates(2); // Initialize with capacity = 2
D.push(1);
D.push(2);
D.push(3);
D.push(4);
D.push(5); // The stacks are now: 2 4
1 3 5
οΉ οΉ οΉ
D.popAtStack(0); // Returns 2. The stacks are now: 4
1 3 5
οΉ οΉ οΉ
D.push(20); // The stacks are now: 20 4
1 3 5
οΉ οΉ οΉ
D.push(21); // The stacks are now: 20 4 21
1 3 5
οΉ οΉ οΉ
D.popAtStack(0); // Returns 20. The stacks are now: 4 21
1 3 5
οΉ οΉ οΉ
D.popAtStack(2); // Returns 21. The stacks are now: 4
1 3 5
οΉ οΉ οΉ
D.pop() // Returns 5. The stacks are now: 4
1 3
οΉ οΉ
D.pop() // Returns 4. The stacks are now: 1 3
οΉ οΉ
D.pop() // Returns 3. The stacks are now: 1
οΉ
D.pop() // Returns 1. There are no stacks.
D.pop() // Returns -1. There are still no stacks.
Constraints:
1 <= capacity <= 2 * 104
1 <= val <= 2 * 104
0 <= index <= 105
At most 2 * 105 calls will be made to push, pop, and popAtStack.
Solutions
Solution 1: Stack Array + Ordered Set
We define the following data structures or variables:
capacity: The capacity of each stack;
stacks: Stack array, used to store all stacks, each with a maximum capacity of capacity;
not_full: Ordered set, used to store the indices of all non-full stacks in the stack array.
For the push(val) operation:
We first check if not_full is empty. If it is, it means there are no non-full stacks, so we need to create a new stack and push val into it. At this point, we check if the capacity capacity is greater than $1$. If it is, we add the index of this stack to not_full.
If not_full is not empty, it means there are non-full stacks. We take out the smallest index index from not_full, and push val into stacks[index]. At this point, if the capacity of stacks[index] equals capacity, we remove index from not_full.
For the popAtStack(index) operation:
We first check if index is within the index range of stacks. If it is not, we directly return $-1$. If stacks[index] is empty, we also directly return $-1$.
If stacks[index] is not empty, we pop the top element val from stacks[index]. If index equals the length of stacks minus $1$, it means stacks[index] is the last stack. If it is empty, we loop to remove the index of the last stack from not_full, and remove the last stack from the stack array stacks, until the last stack is not empty, or the stack array stacks is empty. Otherwise, if stacks[index] is not the last stack, we add index to not_full.
Finally, return val.
For the pop() operation:
We directly call popAtStack(stacks.length - 1).
The time complexity is $(n \times \log n)$, and the space complexity is $O(n)$. Here, $n$ is the number of operations.
fromsortedcontainersimportSortedSetclassDinnerPlates:def__init__(self,capacity:int):self.capacity=capacityself.stacks=[]self.not_full=SortedSet()defpush(self,val:int)->None:ifnotself.not_full:self.stacks.append([val])ifself.capacity>1:self.not_full.add(len(self.stacks)-1)else:index=self.not_full[0]self.stacks[index].append(val)iflen(self.stacks[index])==self.capacity:self.not_full.discard(index)defpop(self)->int:returnself.popAtStack(len(self.stacks)-1)defpopAtStack(self,index:int)->int:ifindex<0orindex>=len(self.stacks)ornotself.stacks[index]:return-1val=self.stacks[index].pop()ifindex==len(self.stacks)-1andnotself.stacks[-1]:whileself.stacksandnotself.stacks[-1]:self.not_full.discard(len(self.stacks)-1)self.stacks.pop()else:self.not_full.add(index)returnval# Your DinnerPlates object will be instantiated and called as such:# obj = DinnerPlates(capacity)# obj.push(val)# param_2 = obj.pop()# param_3 = obj.popAtStack(index)
classDinnerPlates{privateintcapacity;privateList<Deque<Integer>>stacks=newArrayList<>();privateTreeSet<Integer>notFull=newTreeSet<>();publicDinnerPlates(intcapacity){this.capacity=capacity;}publicvoidpush(intval){if(notFull.isEmpty()){stacks.add(newArrayDeque<>());stacks.get(stacks.size()-1).push(val);if(capacity>1){notFull.add(stacks.size()-1);}}else{intindex=notFull.first();stacks.get(index).push(val);if(stacks.get(index).size()==capacity){notFull.pollFirst();}}}publicintpop(){returnpopAtStack(stacks.size()-1);}publicintpopAtStack(intindex){if(index<0||index>=stacks.size()||stacks.get(index).isEmpty()){return-1;}intval=stacks.get(index).pop();if(index==stacks.size()-1&&stacks.get(stacks.size()-1).isEmpty()){while(!stacks.isEmpty()&&stacks.get(stacks.size()-1).isEmpty()){notFull.remove(stacks.size()-1);stacks.remove(stacks.size()-1);}}else{notFull.add(index);}returnval;}}/** * Your DinnerPlates object will be instantiated and called as such: * DinnerPlates obj = new DinnerPlates(capacity); * obj.push(val); * int param_2 = obj.pop(); * int param_3 = obj.popAtStack(index); */
classDinnerPlates{public:DinnerPlates(intcapacity){this->capacity=capacity;}voidpush(intval){if(notFull.empty()){stacks.emplace_back(stack<int>());stacks.back().push(val);if(capacity>1){notFull.insert(stacks.size()-1);}}else{intindex=*notFull.begin();stacks[index].push(val);if(stacks[index].size()==capacity){notFull.erase(index);}}}intpop(){returnpopAtStack(stacks.size()-1);}intpopAtStack(intindex){if(index<0||index>=stacks.size()||stacks[index].empty()){return-1;}intval=stacks[index].top();stacks[index].pop();if(index==stacks.size()-1&&stacks[index].empty()){while(!stacks.empty()&&stacks.back().empty()){notFull.erase(stacks.size()-1);stacks.pop_back();}}else{notFull.insert(index);}returnval;}private:intcapacity;vector<stack<int>>stacks;set<int>notFull;};/** * Your DinnerPlates object will be instantiated and called as such: * DinnerPlates* obj = new DinnerPlates(capacity); * obj->push(val); * int param_2 = obj->pop(); * int param_3 = obj->popAtStack(index); */
typeDinnerPlatesstruct{capacityintstacks[][]intnotFull*redblacktree.Tree}funcConstructor(capacityint)DinnerPlates{returnDinnerPlates{capacity:capacity,notFull:redblacktree.NewWithIntComparator()}}func(this*DinnerPlates)Push(valint){ifthis.notFull.Size()==0{this.stacks=append(this.stacks,[]int{val})ifthis.capacity>1{this.notFull.Put(len(this.stacks)-1,nil)}}else{index,_:=this.notFull.Left().Key.(int)this.stacks[index]=append(this.stacks[index],val)iflen(this.stacks[index])==this.capacity{this.notFull.Remove(index)}}}func(this*DinnerPlates)Pop()int{returnthis.PopAtStack(len(this.stacks)-1)}func(this*DinnerPlates)PopAtStack(indexint)int{ifindex<0||index>=len(this.stacks)||len(this.stacks[index])==0{return-1}val:=this.stacks[index][len(this.stacks[index])-1]this.stacks[index]=this.stacks[index][:len(this.stacks[index])-1]ifindex==len(this.stacks)-1&&len(this.stacks[index])==0{forlen(this.stacks)>0&&len(this.stacks[len(this.stacks)-1])==0{this.notFull.Remove(len(this.stacks)-1)this.stacks=this.stacks[:len(this.stacks)-1]}}else{this.notFull.Put(index,nil)}returnval}/** * Your DinnerPlates object will be instantiated and called as such: * obj := Constructor(capacity); * obj.Push(val); * param_2 := obj.Pop(); * param_3 := obj.PopAtStack(index); */
classDinnerPlates{capacity:number;stacks:number[][];notFull:TreeSet<number>;constructor(capacity:number){this.capacity=capacity;this.stacks=[];this.notFull=newTreeSet<number>();}push(val:number):void{if(this.notFull.size()===0){this.stacks.push([val]);if(this.capacity>1){this.notFull.add(this.stacks.length-1);}}else{constindex=this.notFull.first()!;this.stacks[index].push(val);if(this.stacks[index].length===this.capacity){this.notFull.delete(index);}}}pop():number{returnthis.popAtStack(this.stacks.length-1);}popAtStack(index:number):number{if(index<0||index>=this.stacks.length||this.stacks[index].length===0){return-1;}constval=this.stacks[index].pop()!;if(index===this.stacks.length-1&&this.stacks[index].length===0){while(this.stacks.length>0&&this.stacks[this.stacks.length-1].length===0){this.notFull.delete(this.stacks.length-1);this.stacks.pop();}}else{this.notFull.add(index);}returnval;}}typeCompare<T>=(lhs:T,rhs:T)=>number;classRBTreeNode<T=number>{data:T;count:number;left:RBTreeNode<T>|null;right:RBTreeNode<T>|null;parent:RBTreeNode<T>|null;color:number;constructor(data:T){this.data=data;this.left=this.right=this.parent=null;this.color=0;this.count=1;}sibling():RBTreeNode<T>|null{if(!this.parent)returnnull;// sibling null if no parentreturnthis.isOnLeft()?this.parent.right:this.parent.left;}isOnLeft():boolean{returnthis===this.parent!.left;}hasRedChild():boolean{return(Boolean(this.left&&this.left.color===0)||Boolean(this.right&&this.right.color===0));}}classRBTree<T>{root:RBTreeNode<T>|null;lt:(l:T,r:T)=>boolean;constructor(compare:Compare<T>=(l:T,r:T)=>(l<r?-1:l>r?1:0)){this.root=null;this.lt=(l:T,r:T)=>compare(l,r)<0;}rotateLeft(pt:RBTreeNode<T>):void{constright=pt.right!;pt.right=right.left;if(pt.right)pt.right.parent=pt;right.parent=pt.parent;if(!pt.parent)this.root=right;elseif(pt===pt.parent.left)pt.parent.left=right;elsept.parent.right=right;right.left=pt;pt.parent=right;}rotateRight(pt:RBTreeNode<T>):void{constleft=pt.left!;pt.left=left.right;if(pt.left)pt.left.parent=pt;left.parent=pt.parent;if(!pt.parent)this.root=left;elseif(pt===pt.parent.left)pt.parent.left=left;elsept.parent.right=left;left.right=pt;pt.parent=left;}swapColor(p1:RBTreeNode<T>,p2:RBTreeNode<T>):void{consttmp=p1.color;p1.color=p2.color;p2.color=tmp;}swapData(p1:RBTreeNode<T>,p2:RBTreeNode<T>):void{consttmp=p1.data;p1.data=p2.data;p2.data=tmp;}fixAfterInsert(pt:RBTreeNode<T>):void{letparent=null;letgrandParent=null;while(pt!==this.root&&pt.color!==1&&pt.parent?.color===0){parent=pt.parent;grandParent=pt.parent.parent;/* Case : A Parent of pt is left child of Grand-parent of pt */if(parent===grandParent?.left){constuncle=grandParent.right;/* Case : 1 The uncle of pt is also red Only Recoloring required */if(uncle&&uncle.color===0){grandParent.color=0;parent.color=1;uncle.color=1;pt=grandParent;}else{/* Case : 2 pt is right child of its parent Left-rotation required */if(pt===parent.right){this.rotateLeft(parent);pt=parent;parent=pt.parent;}/* Case : 3 pt is left child of its parent Right-rotation required */this.rotateRight(grandParent);this.swapColor(parent!,grandParent);pt=parent!;}}else{/* Case : B Parent of pt is right child of Grand-parent of pt */constuncle=grandParent!.left;/* Case : 1 The uncle of pt is also red Only Recoloring required */if(uncle!=null&&uncle.color===0){grandParent!.color=0;parent.color=1;uncle.color=1;pt=grandParent!;}else{/* Case : 2 pt is left child of its parent Right-rotation required */if(pt===parent.left){this.rotateRight(parent);pt=parent;parent=pt.parent;}/* Case : 3 pt is right child of its parent Left-rotation required */this.rotateLeft(grandParent!);this.swapColor(parent!,grandParent!);pt=parent!;}}}this.root!.color=1;}delete(val:T):boolean{constnode=this.find(val);if(!node)returnfalse;node.count--;if(!node.count)this.deleteNode(node);returntrue;}deleteAll(val:T):boolean{constnode=this.find(val);if(!node)returnfalse;this.deleteNode(node);returntrue;}deleteNode(v:RBTreeNode<T>):void{constu=BSTreplace(v);// True when u and v are both blackconstuvBlack=(u===null||u.color===1)&&v.color===1;constparent=v.parent!;if(!u){// u is null therefore v is leafif(v===this.root)this.root=null;// v is root, making root nullelse{if(uvBlack){// u and v both black// v is leaf, fix double black at vthis.fixDoubleBlack(v);}else{// u or v is redif(v.sibling()){// sibling is not null, make it red"v.sibling()!.color=0;}}// delete v from the treeif(v.isOnLeft())parent.left=null;elseparent.right=null;}return;}if(!v.left||!v.right){// v has 1 childif(v===this.root){// v is root, assign the value of u to v, and delete uv.data=u.data;v.left=v.right=null;}else{// Detach v from tree and move u upif(v.isOnLeft())parent.left=u;elseparent.right=u;u.parent=parent;if(uvBlack)this.fixDoubleBlack(u);// u and v both black, fix double black at uelseu.color=1;// u or v red, color u black}return;}// v has 2 children, swap data with successor and recursethis.swapData(u,v);this.deleteNode(u);// find node that replaces a deleted node in BSTfunctionBSTreplace(x:RBTreeNode<T>):RBTreeNode<T>|null{// when node have 2 childrenif(x.left&&x.right)returnsuccessor(x.right);// when leafif(!x.left&&!x.right)returnnull;// when single childreturnx.left??x.right;}// find node that do not have a left child// in the subtree of the given nodefunctionsuccessor(x:RBTreeNode<T>):RBTreeNode<T>{lettemp=x;while(temp.left)temp=temp.left;returntemp;}}fixDoubleBlack(x:RBTreeNode<T>):void{if(x===this.root)return;// Reached rootconstsibling=x.sibling();constparent=x.parent!;if(!sibling){// No sibiling, double black pushed upthis.fixDoubleBlack(parent);}else{if(sibling.color===0){// Sibling redparent.color=0;sibling.color=1;if(sibling.isOnLeft())this.rotateRight(parent);// left caseelsethis.rotateLeft(parent);// right casethis.fixDoubleBlack(x);}else{// Sibling blackif(sibling.hasRedChild()){// at least 1 red childrenif(sibling.left&&sibling.left.color===0){if(sibling.isOnLeft()){// left leftsibling.left.color=sibling.color;sibling.color=parent.color;this.rotateRight(parent);}else{// right leftsibling.left.color=parent.color;this.rotateRight(sibling);this.rotateLeft(parent);}}else{if(sibling.isOnLeft()){// left rightsibling.right!.color=parent.color;this.rotateLeft(sibling);this.rotateRight(parent);}else{// right rightsibling.right!.color=sibling.color;sibling.color=parent.color;this.rotateLeft(parent);}}parent.color=1;}else{// 2 black childrensibling.color=0;if(parent.color===1)this.fixDoubleBlack(parent);elseparent.color=1;}}}}insert(data:T):boolean{// search for a position to insertletparent=this.root;while(parent){if(this.lt(data,parent.data)){if(!parent.left)break;elseparent=parent.left;}elseif(this.lt(parent.data,data)){if(!parent.right)break;elseparent=parent.right;}elsebreak;}// insert node into parentconstnode=newRBTreeNode(data);if(!parent)this.root=node;elseif(this.lt(node.data,parent.data))parent.left=node;elseif(this.lt(parent.data,node.data))parent.right=node;else{parent.count++;returnfalse;}node.parent=parent;this.fixAfterInsert(node);returntrue;}find(data:T):RBTreeNode<T>|null{letp=this.root;while(p){if(this.lt(data,p.data)){p=p.left;}elseif(this.lt(p.data,data)){p=p.right;}elsebreak;}returnp??null;}*inOrder(root:RBTreeNode<T>=this.root!):Generator<T,undefined,void>{if(!root)return;for(constvofthis.inOrder(root.left!))yieldv;yieldroot.data;for(constvofthis.inOrder(root.right!))yieldv;}*reverseInOrder(root:RBTreeNode<T>=this.root!):Generator<T,undefined,void>{if(!root)return;for(constvofthis.reverseInOrder(root.right!))yieldv;yieldroot.data;for(constvofthis.reverseInOrder(root.left!))yieldv;}}classTreeSet<T=number>{_size:number;tree:RBTree<T>;compare:Compare<T>;constructor(collection:T[]|Compare<T>=[],compare:Compare<T>=(l:T,r:T)=>(l<r?-1:l>r?1:0),){if(typeofcollection==='function'){compare=collection;collection=[];}this._size=0;this.compare=compare;this.tree=newRBTree(compare);for(constvalofcollection)this.add(val);}size():number{returnthis._size;}has(val:T):boolean{return!!this.tree.find(val);}add(val:T):boolean{constsuccessful=this.tree.insert(val);this._size+=successful?1:0;returnsuccessful;}delete(val:T):boolean{constdeleted=this.tree.deleteAll(val);this._size-=deleted?1:0;returndeleted;}ceil(val:T):T|undefined{letp=this.tree.root;lethigher=null;while(p){if(this.compare(p.data,val)>=0){higher=p;p=p.left;}else{p=p.right;}}returnhigher?.data;}floor(val:T):T|undefined{letp=this.tree.root;letlower=null;while(p){if(this.compare(val,p.data)>=0){lower=p;p=p.right;}else{p=p.left;}}returnlower?.data;}higher(val:T):T|undefined{letp=this.tree.root;lethigher=null;while(p){if(this.compare(val,p.data)<0){higher=p;p=p.left;}else{p=p.right;}}returnhigher?.data;}lower(val:T):T|undefined{letp=this.tree.root;letlower=null;while(p){if(this.compare(p.data,val)<0){lower=p;p=p.right;}else{p=p.left;}}returnlower?.data;}first():T|undefined{returnthis.tree.inOrder().next().value;}last():T|undefined{returnthis.tree.reverseInOrder().next().value;}shift():T|undefined{constfirst=this.first();if(first===undefined)returnundefined;this.delete(first);returnfirst;}pop():T|undefined{constlast=this.last();if(last===undefined)returnundefined;this.delete(last);returnlast;}*[Symbol.iterator]():Generator<T,void,void>{for(constvalofthis.values())yieldval;}*keys():Generator<T,void,void>{for(constvalofthis.values())yieldval;}*values():Generator<T,undefined,void>{for(constvalofthis.tree.inOrder())yieldval;returnundefined;}/** * Return a generator for reverse order traversing the set */*rvalues():Generator<T,undefined,void>{for(constvalofthis.tree.reverseInOrder())yieldval;returnundefined;}}/** * Your DinnerPlates object will be instantiated and called as such: * var obj = new DinnerPlates(capacity) * obj.push(val) * var param_2 = obj.pop() * var param_3 = obj.popAtStack(index) */