Problem

Design a stack that supports increment operations on its elements.

Implement the CustomStack class:

  • CustomStack(int maxSize) Initializes the object with maxSize which is the maximum number of elements in the stack.
  • void push(int x) Adds x to the top of the stack if the stack has not reached the maxSize.
  • int pop() Pops and returns the top of the stack or -1 if the stack is empty.
  • void inc(int k, int val) Increments the bottom k elements of the stack by val. If there are less than k elements in the stack, increment all the elements in the stack.

Examples

Example 1:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

Input
["CustomStack","push","push","pop","push","push","push","increment","increment","pop","pop","pop","pop"]
[[3],[1],[2],[],[2],[3],[4],[5,100],[2,100],[],[],[],[]]

Output
[null,null,null,2,null,null,null,null,null,103,202,201,-1]

Explanation
CustomStack stk = new CustomStack(3); // Stack is Empty []
stk.push(1);                          // stack becomes [1]
stk.push(2);                          // stack becomes [1, 2]
stk.pop();                            // return 2 --> Return top of the stack 2, stack becomes [1]
stk.push(2);                          // stack becomes [1, 2]
stk.push(3);                          // stack becomes [1, 2, 3]
stk.push(4);                          // stack still [1, 2, 3], Do not add another elements as size is 4
stk.increment(5, 100);                // stack becomes [101, 102, 103]
stk.increment(2, 100);                // stack becomes [201, 202, 103]
stk.pop();                            // return 103 --> Return top of the stack 103, stack becomes [201, 202]
stk.pop();                            // return 202 --> Return top of the stack 202, stack becomes [201]
stk.pop();                            // return 201 --> Return top of the stack 201, stack becomes []
stk.pop();                            // return -1 --> Stack is empty return -1.

Solution

Method 1 - Using an array

We can use array based stack implementation, and handle the increments just by looping till k initial elements.

Code

 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
37
38
39

class CustomStack {
    private int maxSize;
    private int[] stack;
    private int top;

    public CustomStack(int maxSize) {
        this.maxSize = maxSize;
        this.stack = new int[maxSize];
        this.top = 0;
    }

    public void push(int x) {
        if (top < maxSize) {
            stack[top++] = x;
        }
    }

    public int pop() {
        if (top == 0) {
            return -1;
        }
        return stack[--top];
    }

    public void increment(int k, int val) {
        int limit = Math.min(k, top);
        for (int i = 0; i < limit; i++) {
            stack[i] += val;
        }
    }
}

/**
 * Your CustomStack object will be instantiated and called as such:
 * CustomStack obj = new CustomStack(maxSize);
 * obj.push(x);
 * int param_2 = obj.pop();
 * obj.increment(k,val);
 */
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

class CustomStack:

    def __init__(self, maxSize: int):
        self.stack = []
        self.maxSize = maxSize

    def push(self, x: int) -> None:
        if len(self.stack) < self.maxSize:
            self.stack.append(x)

    def pop(self) -> int:
        return self.stack.pop() if self.stack else -1

    def increment(self, k: int, val: int) -> None:
        limit = min(k, len(self.stack))
        for i in range(limit):
            self.stack[i] += val

Complexity

  • ⏰ Time complexity
    • push(x): O(1) - Adding an element to the stack.
    • pop(): O(1) - Removing an element from the stack.
    • increment(k, val): O(k) - Incrementing the value of the bottom k elements.
  • 🧺 Space complexity: O(n), , where n is the maxSize of the stack. This accounts for the space used by the single stack array.