Skip to content

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()}|)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
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;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
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<>();
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
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