Design a Stack With Increment Operation

Med

#1280Time: - **`push()`**: O(1) - Constant time array access. - **`pop()`**: O(1) - Constant time array access. - **`increment()`**: O(k) - In the worst case, we iterate through `k` elements, which can be up to `maxSize`. So, it's O(maxSize).Space: O(maxSize) - We need an array of size `maxSize` to store the stack elements.4 companies

Patterns

Design

Data structures

Array Stack

Companies

eBay Cloudflare Moloco IMC

Prompt

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.

Example 1:

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.

Constraints:

1 <= maxSize, x, k <= 1000
0 <= val <= 100
At most 1000 calls will be made to each method of increment, push and pop each separately.

Approaches

2 approaches with complexity analysis and trade-offs.

This approach uses a simple array or list to simulate the stack. The push and pop operations work as they would in a standard array-based stack implementation. The increment operation is handled by directly iterating through the bottom k elements of the array and adding the value val to each one.

Algorithm

Initialize an integer array stack of size maxSize to store the elements.
Maintain an integer top to point to the index of the top element, initialized to -1.
push(x): If the stack is not full (top < maxSize - 1), increment top and place x at stack[top].
pop(): If the stack is not empty (top >= 0), return the element at stack[top] and then decrement top. Otherwise, return -1.
inc(k, val): Iterate from index i = 0 up to min(k - 1, top). In each iteration, add val to stack[i]. This directly modifies the values of the bottom elements.

Walkthrough

In this method, we use a standard array to represent the stack's data. A pointer, top, keeps track of the current size and the position for the next push.

The push and pop methods are implemented in the standard O(1) time complexity for an array-based stack.
The increment(k, val) method is the most straightforward part. It loops through the first k elements of the array (or fewer if the stack size is less than k) and adds val to each. This makes the increment operation's performance dependent on the value of k.

1class CustomStack {2    private int[] stack;3    private int top;4    private int maxSize;5 6    public CustomStack(int maxSize) {7        this.maxSize = maxSize;8        this.stack = new int[maxSize];9        this.top = -1;10    }11 12    public void push(int x) {13        if (top < maxSize - 1) {14            top++;15            stack[top] = x;16        }17    }18 19    public int pop() {20        if (top == -1) {21            return -1;22        }23        int val = stack[top];24        top--;25        return val;26    }27 28    public void increment(int k, int val) {29        int limit = Math.min(k, top + 1);30        for (int i = 0; i < limit; i++) {31            stack[i] += val;32        }33    }34}

Complexity

Time

- **`push()`**: O(1) - Constant time array access. - **`pop()`**: O(1) - Constant time array access. - **`increment()`**: O(k) - In the worst case, we iterate through `k` elements, which can be up to `maxSize`. So, it's O(maxSize).

Space

O(maxSize) - We need an array of size `maxSize` to store the stack elements.

Trade-offs

Pros

Simple to understand and implement.
Uses a standard array, which is memory-efficient.

Cons

The increment operation is inefficient, with a time complexity linear to k. This can be slow if k is large and increment is called frequently.

This approach avoids the costly iteration in the increment method by deferring the additions. It uses an auxiliary array to store the increments, making the increment operation O(1). The actual value of an element is calculated only when it is popped from the stack.

Algorithm

Initialize a value array stack, an auxiliary increment array inc of the same size, and a top pointer.
push(x): If not full, increment top, set stack[top] = x, and ensure inc[top] is 0.
increment(k, val): If the stack is not empty, find the index of the top-most element to be incremented, idx = min(k-1, top). Add val to inc[idx]. This is an O(1) step.
pop(): If not empty, get the total increment for the top element: increment_value = inc[top]. Calculate the final value: result = stack[top] + increment_value. Before popping, pass the increment down to the element below: inc[top-1] += increment_value (if top > 0). Reset inc[top] = 0, decrement top, and return result.

Walkthrough

To optimize the increment operation, we can avoid iterating through the elements. The core idea is to use a lazy update mechanism. We maintain a second array, let's call it inc, of the same size as the stack.

inc[i] will store the total increment value that needs to be applied to stack[i] and all elements below it.
When increment(k, val) is called, instead of adding val to k elements, we simply add val to inc[k-1] (or the top-most valid index). This marks that the element at this position and everything below it has an additional val to be added.
When pop() is called for the element at top, we calculate its true value by adding stack[top] + inc[top]. Then, we propagate the increment downwards by adding inc[top] to inc[top-1]. This ensures the element below will correctly account for the increment when it is popped. This makes all operations O(1).

1class CustomStack {2    private int[] stack;3    private int[] inc; // Stores lazy increments4    private int top;5    private int maxSize;6 7    public CustomStack(int maxSize) {8        this.maxSize = maxSize;9        this.stack = new int[maxSize];10        this.inc = new int[maxSize];11        this.top = -1;12    }13 14    public void push(int x) {15        if (top < maxSize - 1) {16            top++;17            stack[top] = x;18            inc[top] = 0; // Reset increment for new element19        }20    }21 22    public int pop() {23        if (top == -1) {24            return -1;25        }26        27        int incrementValue = inc[top];28        int result = stack[top] + incrementValue;29        30        // Pass the increment down to the element below31        if (top > 0) {32            inc[top - 1] += incrementValue;33        }34        35        top--;36        37        return result;38    }39 40    public void increment(int k, int val) {41        if (top == -1) {42            return;43        }44        // The increment applies to elements from index 0 to min(k, size)-1.45        // We only need to add the increment to the top-most affected element's slot.46        int limitIndex = Math.min(k - 1, top);47        inc[limitIndex] += val;48    }49}

Complexity

Time

- **`push()`**: O(1) - Constant time. - **`pop()`**: O(1) - Constant time. - **`increment()`**: O(1) - Constant time. All operations are now highly efficient.

Space

O(maxSize) - We use two arrays, `stack` and `inc`, both of size `maxSize`. The total space is `O(2 * maxSize)`, which simplifies to O(maxSize).

Trade-offs

Pros

Extremely efficient, with all operations taking constant time.
Scales well regardless of the number of calls or the value of k.

Cons

Requires extra space for the increment array, doubling the storage requirement compared to the naive approach.
The logic is slightly more complex to understand due to the deferred nature of the increment operation.

Solutions

Solution

1class CustomStack { private int [] stk ; private int [] add ; private int i ; public CustomStack ( int maxSize ) { stk = new int [ maxSize ]; add = new int [ maxSize ]; } public void push ( int x ) { if ( i < stk . length ) { stk [ i ++] = x ; } } public int pop () { if ( i <= 0 ) { return - 1 ; } int ans = stk [-- i ] + add [ i ]; if ( i > 0 ) { add [ i - 1 ] += add [ i ]; } add [ i ] = 0 ; return ans ; } public void increment ( int k , int val ) { if ( i > 0 ) { add [ Math . min ( i , k ) - 1 ] += 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); */

1class CustomStack { public: CustomStack ( int maxSize ) { stk . resize ( maxSize ); add . resize ( maxSize ); i = 0 ; } void push ( int x ) { if ( i < stk . size ()) { stk [ i ++ ] = x ; } } int pop () { if ( i <= 0 ) { return - 1 ; } int ans = stk [ -- i ] + add [ i ]; if ( i > 0 ) { add [ i - 1 ] += add [ i ]; } add [ i ] = 0 ; return ans ; } void increment ( int k , int val ) { if ( i > 0 ) { add [ min ( k , i ) - 1 ] += val ; } } private: vector < int > stk ; vector < int > add ; int i ; }; /** * 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); */

1class CustomStack : def __init__ ( self , maxSize : int ): self . stack = [] self . max_size = maxSize self . inc = [] def push ( self , x : int ) -> None : if len ( self . stack ) < self . max_size : self . stack . append ( x ) self . inc . append ( 0 ) # 0 because no range operation # not append(inc[-1]), since it will be updated on pop() def pop ( self ) -> int : if not self . stack : return - 1 if len ( self . inc ) > 1 : self . inc [ - 2 ] += self . inc [ - 1 ] return self . inc . pop () + self . stack . pop () def increment ( self , k : int , val : int ) -> None : if self . inc : self . inc [ min ( k , len ( self . inc )) - 1 ] += val # Your CustomStack object will be instantiated and called as such: # obj = CustomStack(maxSize) # obj.push(x) # param_2 = obj.pop() # obj.increment(k,val) ############### class CustomStack : def __init__ ( self , maxSize : int ): self . stk = [ 0 ] * maxSize self . add = [ 0 ] * maxSize self . i = 0 def push ( self , x : int ) -> None : if self . i < len ( self . stk ): self . stk [ self . i ] = x self . i += 1 def pop ( self ) -> int : if self . i <= 0 : return - 1 self . i -= 1 ans = self . stk [ self . i ] + self . add [ self . i ] if self . i > 0 : self . add [ self . i - 1 ] += self . add [ self . i ] self . add [ self . i ] = 0 return ans def increment ( self , k : int , val : int ) -> None : i = min ( k , self . i ) - 1 if i >= 0 : self . add [ i ] += val # Your CustomStack object will be instantiated and called as such: # obj = CustomStack(maxSize) # obj.push(x) # param_2 = obj.pop() # obj.increment(k,val)

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem