Formatted question description: https://leetcode.ca/all/1172.html

# 1172. Dinner Plate Stacks

Hard

## Description

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.
• void push(int val) pushes the given positive integer val into the leftmost stack with 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 stacks are empty.
• int popAtStack(int index) returns the value at the top of the stack with the given index and removes it from that stack, and returns -1 if the stack with that given index is empty.

Example:

Input:
["DinnerPlates","push","push","push","push","push","popAtStack","push","push","popAtStack","popAtStack","pop","pop","pop","pop","pop"]
[,,,,,,,,,,,[],[],[],[],[]]
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 <= 20000
• 1 <= val <= 20000
• 0 <= index <= 100000
• At most 200000 calls will be made to push, pop, and popAtStack.

## Solution

In the class, maintain capacity that is the capacity of each stack, a tree map that stores each index and the corresponding stack, a tree set that stores the indices of available stacks, and the current stack’s index.

For the constructor, initialize the fields, where the current stack’s index is initialized to 0.

For method push, check whether there is any element in the tree set. If so, obtain the minimum index from the tree set and the corresponding stack and add the element val to the stack, and remove the minimum index from the tree set if the stack becomes full. Otherwise, add the element val to the stack at the current stack’s index, and increase the current stack’s index by 1 if the stack becomes full.

For method pop, if the tree map does not have any entry, return -1. Otherwise, obtain the maximum index in the tree map and obtain the corresponding stack, and pop one element from the stack. Update the current stack’s index if the last stack is not full. Finally, return the popped element.

For method popAtIndex, first obtain the stack at the given index. If the stack is empty, return -1. Otherwise, pop one element from the stack, and update the current stack’s index if necessary or add the given index to the tree set if the given index is not equal to or adjacent to the current stack’s index. Finally, return the popped element.

class DinnerPlates {
int capacity;
TreeMap<Integer, Stack<Integer>> stackMap;
TreeSet<Integer> availableStacks;
int stackIndex;

public DinnerPlates(int capacity) {
this.capacity = capacity;
stackMap = new TreeMap<Integer, Stack<Integer>>();
availableStacks = new TreeSet<Integer>();
stackIndex = 0;
}

public void push(int val) {
if (availableStacks.isEmpty()) {
Stack<Integer> stack = stackMap.getOrDefault(stackIndex, new Stack<Integer>());
stack.push(val);
stackMap.put(stackIndex, stack);
if (stack.size() == capacity)
stackIndex++;
} else {
int index = availableStacks.first();
Stack<Integer> stack = stackMap.getOrDefault(index, new Stack<Integer>());
stack.push(val);
stackMap.put(index, stack);
if (stack.size() == capacity)
availableStacks.remove(index);
}
}

public int pop() {
if (stackMap.size() == 0)
return -1;
int lastKey = stackMap.lastKey();
Stack<Integer> stack = stackMap.get(lastKey);
int element = stack.pop();
if (stack.isEmpty())
stackMap.remove(lastKey);
if (stackIndex - lastKey == 1)
stackIndex = lastKey;
return element;
}

public int popAtStack(int index) {
Stack<Integer> stack = stackMap.getOrDefault(index, new Stack<Integer>());
if (stack.isEmpty())
return -1;
else {
int element = stack.pop();
if (stack.isEmpty())
stackMap.remove(index);
if (stackIndex - index == 1)
stackIndex = index;
else if (stackIndex - index > 1)
return element;
}
}
}

/**
* 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);
*/