Dinner Plate Stacks

To overcome the performance issues of the naive approach, we can use more suitable data structures. A TreeMap can store the stacks, and a TreeSet can track which stacks have available capacity. This allows us to find the target stacks for push and pop operations in logarithmic time instead of linear time.

• Data Structures:
  • TreeMap<Integer, Deque<Integer>> stacks: A map from a stack's index to the stack itself (using a Deque). TreeMap keeps indices sorted, allowing efficient retrieval of the rightmost stack.
  • TreeSet<Integer> availableStacks: A sorted set of indices of stacks that are not full. This allows efficient retrieval of the leftmost available stack.
  • int capacity: The maximum capacity of each stack.
• Constructor DinnerPlates(capacity):
  • Initialize capacity, the TreeMap, and the TreeSet.
• push(val) Operation:
  1. Check if availableStacks is empty. If so, it means all existing stacks are full. Determine the index for a new stack (which will be stacks.lastKey() + 1, or 0 if no stacks exist) and add this index to availableStacks.
  2. Get the smallest index from availableStacks using first(). This is the leftmost stack with available space.
  3. Get the stack at this index from the stacks map (creating it if it doesn't exist) and push val onto it.
  4. If the stack's size now equals capacity, remove its index from availableStacks.
• pop() Operation:
  1. Check if the stacks map is empty. If so, return -1.
  2. Find the index of the rightmost non-empty stack using stacks.lastKey().
  3. Delegate the actual pop logic to popAtStack(index) to ensure consistency.
• popAtStack(index) Operation:
  1. Check if a stack exists at the given index and if it's not empty. If not, return -1.
  2. Note whether the stack was full before the pop operation.
  3. Pop the value from the stack.
  4. If the stack becomes empty, remove it entirely from the stacks map and also from availableStacks.
  5. If the stack was full before popping, it now has space, so add its index back to availableStacks.
  6. Return the popped value.

Data Structures

This optimized solution uses two primary data structures to manage the stacks efficiently:

1. TreeMap<Integer, Deque<Integer>> stacks: This maps a stack's index to the stack itself (an ArrayDeque is used as the stack implementation). A TreeMap is chosen because it keeps its keys (the stack indices) in sorted order. This property is crucial for the pop() operation, as we can find the rightmost non-empty stack in O(log N) time by calling stacks.lastKey().
2. TreeSet<Integer> availableStacks: This stores the indices of all stacks that currently have a size less than capacity. A TreeSet maintains these indices in sorted order, which is perfect for the push(val) operation. We can find the leftmost stack with available space in O(log K) time by calling availableStacks.first(), where K is the number of non-full stacks.

Algorithm Details

• push(val): Instead of scanning, we directly query availableStacks. If it's not empty, first() gives us the index of the leftmost stack with space. We push to it. If it becomes full, we remove its index from availableStacks. If availableStacks is empty, we know we need a new stack at the end of the current row, so we create one at index stacks.lastKey() + 1.

• pop(): We find the rightmost non-empty stack by getting the largest key in our stacks TreeMap (lastKey()). Then, we simply call popAtStack on that index. This avoids a linear scan.

• popAtStack(index): We pop from the specified stack. The key is to correctly update our data structures. If a pop makes a stack empty, we remove it from the stacks map. If a pop makes a full stack become not full, we add its index to availableStacks so it can be used for future pushes.

Code Snippet

import java.util.*;

class DinnerPlates {
    private int capacity;
    private TreeMap<Integer, Deque<Integer>> stacks;
    private TreeSet<Integer> availableStacks;

    public DinnerPlates(int capacity) {
        this.capacity = capacity;
        this.stacks = new TreeMap<>();
        this.availableStacks = new TreeSet<>();
    }

    public void push(int val) {
        if (availableStacks.isEmpty()) {
            int newIndex = stacks.isEmpty() ? 0 : stacks.lastKey() + 1;
            availableStacks.add(newIndex);
        }
        
        int indexToPush = availableStacks.first();
        stacks.computeIfAbsent(indexToPush, k -> new ArrayDeque<>()).push(val);
        
        if (stacks.get(indexToPush).size() == capacity) {
            availableStacks.remove(indexToPush);
        }
    }

    public int pop() {
        if (stacks.isEmpty()) {
            return -1;
        }
        int indexToPop = stacks.lastKey();
        return popAtStack(indexToPop);
    }

    public int popAtStack(int index) {
        if (!stacks.containsKey(index) || stacks.get(index).isEmpty()) {
            return -1;
        }
        
        Deque<Integer> stack = stacks.get(index);
        boolean wasFull = (stack.size() == capacity);
        
        int val = stack.pop();
        
        if (wasFull) {
            availableStacks.add(index);
        }
        
        if (stack.isEmpty()) {
            stacks.remove(index);
            availableStacks.remove(index);
        }
        
        return val;
    }
}