1172. Dinner Plate Stacks

• Time: Constructor: $O(1)$, push(val: int): $O(\log n)$, pop(): $O(\log n)$, popAtStack(index: int): $O(\log n)$
• Space: $O(|\texttt{push()}|)$

C++JavaPython

class DinnerPlates {
 public:
  DinnerPlates(int capacity) : capacity(capacity) {
    minHeap.push(0);
  }

  void push(int val) {
    const int index = minHeap.top();
    // Add new stack on demand
    if (index == stacks.size())
      stacks.push_back({});
    // Push the new value
    stacks[index].push(val);
    // If the stack pushed is full, remove its candidacy from minHeap
    if (stacks[index].size() == capacity) {
      minHeap.pop();
      // If minHeap is empty, the next available stack index is stacks.size()
      if (minHeap.empty())
        minHeap.push(stacks.size());
    }
  }

  int pop() {
    // Remove empty stacks from back
    while (!stacks.empty() && stacks.back().empty())
      stacks.pop_back();
    if (stacks.empty())
      return -1;
    return popAtStack(stacks.size() - 1);
  }

  int popAtStack(int index) {
    if (index >= stacks.size() || stacks[index].empty())
      return -1;
    // If the stack is going to have space, add its candiday to minHeap
    if (stacks[index].size() == capacity)
      minHeap.push(index);
    const int val = stacks[index].top();
    stacks[index].pop();
    return val;
  }

 private:
  const int capacity;
  vector<stack<int>> stacks;
  // Min stacks to push
  priority_queue<int, vector<int>, greater<>> minHeap;
};

class DinnerPlates {
  public DinnerPlates(int capacity) {
    this.capacity = capacity;
    minHeap.offer(0);
  }

  public void push(int val) {
    final int index = minHeap.peek();
    // Add new stack on demand
    if (index == stacks.size())
      stacks.add(new ArrayDeque<>());
    // Push the new value
    stacks.get(index).push(val);
    // If the stack pushed is full, remove its candidacy from minHeap
    if (stacks.get(index).size() == capacity) {
      minHeap.poll();
      // If minHeap is empty, the next available stack index is stacks.size()
      if (minHeap.isEmpty())
        minHeap.offer(stacks.size());
    }
  }

  public int pop() {
    // Remove empty stacks from back
    while (!stacks.isEmpty() && stacks.get(stacks.size() - 1).isEmpty())
      stacks.remove(stacks.size() - 1);
    if (stacks.isEmpty())
      return -1;
    return popAtStack(stacks.size() - 1);
  }

  public int popAtStack(int index) {
    if (index >= stacks.size() || stacks.get(index).isEmpty())
      return -1;
    // If the stack is going to have space, add its candiday to minHeap
    if (stacks.get(index).size() == capacity)
      minHeap.offer(index);
    return stacks.get(index).pop();
  }

  private final int capacity;
  private List<Deque<Integer>> stacks = new ArrayList<>();
  private Queue<Integer> minHeap = new PriorityQueue<>();
}

class DinnerPlates:
  def __init__(self, capacity: int):
    self.capacity = capacity
    self.stacks = []
    self.minHeap = [0]  # Min stacks to push

  def push(self, val: int) -> None:
    index = self.minHeap[0]
    # Add new stack on demand
    if index == len(self.stacks):
      self.stacks.append([])
    # Push the new value
    self.stacks[index].append(val)
    # If the stack pushed is full, remove its candidacy from minHeap
    if len(self.stacks[index]) == self.capacity:
      heapq.heappop(self.minHeap)
      # If minHeap is empty, the next available stack index is stacks.size()
      if not self.minHeap:
        heapq.heappush(self.minHeap, len(self.stacks))

  def pop(self) -> int:
    # Remove empty stacks from back
    while self.stacks and not self.stacks[-1]:
      self.stacks.pop()
    if not self.stacks:
      return -1
    return self.popAtStack(len(self.stacks) - 1)

  def popAtStack(self, index: int) -> int:
    if index >= len(self.stacks) or not self.stacks[index]:
      return -1
    # If the stack is going to have space, add its candiday to minHeap
    if len(self.stacks[index]) == self.capacity:
      heapq.heappush(self.minHeap, index)
    return self.stacks[index].pop()