1381. Design a Stack With Increment Operation

• Time:
  • Constructor: $O(1)$ (C++/Python) | $O(\texttt{maxSize})$ (Java)
  • push(x: int): $O(1)$
  • pop(): $O(1)$
  • increment(k: int, x: val): $O(1)$
• Space: $O(|\texttt{push()}|)$

C++JavaPython

class CustomStack {
 public:
  CustomStack(int maxSize) : maxSize(maxSize) {}

  void push(int x) {
    if (stack.size() == maxSize)
      return;
    stack.push(x);
    pendingIncrements.push_back(0);
  }

  int pop() {
    if (stack.empty())
      return -1;
    const int i = stack.size() - 1;
    if (i > 0)
      pendingIncrements[i - 1] += pendingIncrements[i];
    const int val = stack.top() + pendingIncrements[i];
    stack.pop();
    pendingIncrements.pop_back();
    return val;
  }

  void increment(int k, int val) {
    if (stack.empty())
      return;
    const int i = min(k - 1, static_cast<int>(stack.size()) - 1);
    pendingIncrements[i] += val;
  }

 private:
  const int maxSize;
  stack<int> stack;
  // pendingIncrements[i] := the pending increment for stack[0..i].
  vector<int> pendingIncrements;
};

class CustomStack {
  public CustomStack(int maxSize) {
    this.maxSize = maxSize;
  }

  public void push(int x) {
    if (stack.size() == maxSize)
      return;
    stack.push(x);
    pendingIncrements.add(0);
  }

  public int pop() {
    if (stack.isEmpty())
      return -1;
    final int i = stack.size() - 1;
    final int pendingIncrement = pendingIncrements.get(i);
    pendingIncrements.remove(i);
    if (i > 0)
      pendingIncrements.set(i - 1, pendingIncrements.get(i - 1) + pendingIncrement);
    return stack.pop() + pendingIncrement;
  }

  public void increment(int k, int val) {
    if (stack.isEmpty())
      return;
    final int i = Math.min(k - 1, stack.size() - 1);
    pendingIncrements.set(i, pendingIncrements.get(i) + val);
  }

  private int maxSize;
  private Deque<Integer> stack = new ArrayDeque<>();
  // pendingIncrements[i] := the pending increment for stack[0..i].
  private List<Integer> pendingIncrements = new ArrayList<>();
}

class CustomStack:
  def __init__(self, maxSize: int):
    self.maxSize = maxSize
    self.stack = []
    # pendingIncrements[i] := the pending increment for stack[0..i].
    self.pendingIncrements = []

  def push(self, x: int) -> None:
    if len(self.stack) == self.maxSize:
      return
    self.stack.append(x)
    self.pendingIncrements.append(0)

  def pop(self) -> int:
    if not self.stack:
      return -1
    if len(self.stack) > 1:
      self.pendingIncrements[-2] += self.pendingIncrements[-1]
    return self.stack.pop() + self.pendingIncrements.pop()

  def increment(self, k: int, val: int) -> None:
    if not self.stack:
      return
    i = min(k - 1, len(self.stack) - 1)
    self.pendingIncrements[i] += val