Problem Description

In this problem, we are given a max capacity for stacks and provided with methods to push a new element to these stacks, pop an element from any specific stack, and pop an element from all of these stacks.

• push(int val): This function pushes the given integer value into the leftmost stack that can accommodate more elements up to its maximum capacity.
• pop(): This function returns the top-most element from the rightmost non-empty stack.
• popAtStack(int index): This function returns the top-most element from a specified stack.

For a better understanding, let's walk through the given example:

We are given capacity=2, which means each stack can hold 2 elements at max. Initially, all stacks are empty.

push(1), push(2), push(3), push(4), push(5) Here, first, we push 1 in the first stack, then 2 in the same stack as it can still hold one more element. Now, 3 is pushed in the second stack, then 4 again in the second stack. Finally, the 5 is pushed into the third stack. So, stacks are now 1 3 5↧ 2 4 ↧.

popAtStack(0) Now we pop the last element from stack 0. The 2 is removed. The stacks are now 1 3 5↧ 4 ↧.

Similarly, after several operations, we reach a condition where all stacks are empty. If we try to pop an element, we get -1 which shows all stacks are empty.

Solution Explanation

We use a vector of stacks and a min-heap to solve this problem.

The vector of stacks keeps all the stacks, and the min-heap is responsible for finding the leftmost non-empty stack at O(logn). We keep pushing the index of the non-full stack to min-heap, which is used to find the leftmost stack that is not full. If the top of the min-heap points to a full stack, we pop this index out and push the new leftmost non-full stack index if possible.

To implement popAtStack operation, we directly pop elements from the specified stack and check whether this stack is non-full; if so, we push the index of this stack into the min-heap.

We implement the pop operation by popping elements from the rightmost non-empty stack, and keep deleting the empty stack from the right end of the vector of stacks.

Python Solution

1
2python
3import heapq
4class DinnerPlates:
5
6    def __init__(self, capacity: int):
7        self.stack = []
8        self.capacity = capacity
9        self.available = []
10
11    def push(self, val: int) -> None:
12        while self.available and self.available[0] < len(self.stack) and len(self.stack[self.available[0]]) == self.capacity:
13            heapq.heappop(self.available)
14        if not self.available or self.available[0] == len(self.stack):
15            self.stack.append([])
16            heapq.heappush(self.available, len(self.stack)-1)
17        self.stack[self.available[0]].append(val)
18
19    def pop(self) -> int:
20        while self.stack and not self.stack[-1]:
21            self.stack.pop()
22        return self.popAtStack(len(self.stack)-1)
23
24    def popAtStack(self, index: int) -> int:
25        if 0 <= index < len(self.stack) and self.stack[index]:
26            heapq.heappush(self.available, index)
27            return self.stack[index].pop()
28        return -1

Java Solution

1
2java
3class DinnerPlates {
4  private static final int MAX_CAPACITY = 20001;
5  private int capacity;
6  private Map<Integer, Deque<Integer>> stackMap;
7  private TreeSet<Integer> availableStacks;
8  private int lastStack;
9  
10  public DinnerPlates(int capacity) {
11    this.capacity = capacity;
12    this.stackMap = new HashMap<>();
13    this.availableStacks = new TreeSet<>();
14    this.lastStack = 0;
15  }
16  
17  public void push(int val) {
18    while (!availableStacks.isEmpty() && stackMap.get(availableStacks.first()).size() == capacity) {
19      availableStacks.remove(availableStacks.first());
20    }
21    int stackIdx = availableStacks.isEmpty() ? lastStack : availableStacks.first();
22    stackMap.computeIfAbsent(stackIdx, k -> new ArrayDeque<>()).push(val);
23    availableStacks.add(stackIdx + 1);
24    lastStack = Math.max(lastStack, stackIdx + 1);
25  }
26  
27  public int pop() {
28    if (lastStack == 0) return -1;
29    
30    int val = popAtStack(lastStack - 1);
31    while (lastStack > 0 && stackMap.get(lastStack - 1).isEmpty()) {
32      stackMap.remove(--lastStack);
33    }
34    return val;
35  }
36  
37  public int popAtStack(int index) {
38    Deque<Integer> stack = stackMap.get(index);
39    if (stack == null || stack.isEmpty()) return -1;
40    availableStacks.add(index);
41    return stack.pop();
42  }
43}

JavaScript Solution

1
2javascript
3class DinnerPlates {
4    constructor(capacity) {
5        this.capacity = capacity;
6        this.stacks = [];
7        this.front = new MinPriorityQueue();
8        this.back = new MinPriorityQueue();
9    }
10    
11    push(val) {
12        while (!this.front.isEmpty() && this.stacks[this.front.peek().val].length === this.capacity)
13            this.front.poll();
14        if (this.front.isEmpty() || this.front.peek().val === this.stacks.length)
15            this.stacks.push([]);
16        this.stacks[this.front.peek().val].push(val);
17        if (this.stacks[this.front.peek().val].length === this.capacity)
18            this.front.poll();
19        if (this.front.isEmpty())
20            this.front.enqueue(this.stacks.length, this.stacks.length);
21    }
22    
23    pop() {
24        while (!this.back.isEmpty() && (this.stacks[this.back.peek().val] === undefined || this.stacks[this.back.peek().val].length === 0)) {
25            this.back.poll();
26            this.front.poll();
27        }
28        if (this.back.isEmpty())
29            return -1;
30        const val = this.stacks[this.back.peek().val].pop()
31        this.front.enqueue(this.back.peek().val, this.back.peek().val);
32        return val;
33    }
34    
35    popAtStack(index) {
36        if (this.stacks.length <= index || this.stacks[index].length === 0)
37            return -1;
38        const val = this.stacks[index].pop();
39        this.front.enqueue(index, index);
40        while (!this.back.isEmpty() && (this.stacks[this.back.peek().val] === undefined || this.stacks[this.back.peek().val].length === 0)) {
41            this.back.poll();
42            this.front.poll();
43        }
44        return val;
45    }
46}

+1

The primary component in these solutions is the implementation and use of a min-heap, and a vector (Python) or map (Java) of stacks. In all the three solutions, we place the stacks side by side in a vector (array) or map. Whenever we need to push an element, we check the leftmost stack that can accommodate an element, and push the element to that stack. We use a min-heap to find this stack effectively.

Whenever we need to pop an element, we always pop from the rightmost stack, that has elements, in Python and Java implementations. However, in the JavaScript solution, we pop the rightmost stack with elements or the stack introduced to accommodate popped elements in previous pop operations, whichever is to the extreme right.

The popAtStack(index) operation is pretty straightforward in all the solutions where we pop the index th stack if it has elements. If the pop operation makes room for more elements in the stack, we push the index into the min-heap.

Python's heapq module provides an implementation of the heap queue algorithm, also known as the priority queue algorithm. In this solution, we keep track of the available stacks in a heap and pop the least index only when necessary.

In Java solution, we use Java's in-built data structures TreeSet and Map to store the stacks and track available stacks. TreeSet is a NavigableSet implementation, which allows retrieving elements by their absolute positions.

JavaScript's version requires implementing a MinPriorityQueue() class first, which is a flexible data structure allowing insertion of elements with arbitrary priorities and retrieval of elements in minimum-priority order. This is used in conjunction with an array of stacks.

All solutions essentially maintain state about which stacks are available to push new elements into, and which are available to pop elements from, which minimizes the need to loop through the stacks each time we want to add or remove an element.

The solutions above provide an efficient way for pushing, popping and popping at a specific stack, in a multi-stack with a max capacity data structure by leveraging different data structures and built-in modules in Python, Java, and JavaScript. They handle all the edge cases, making them robust for any possible input.