Design a Stack With Increment Operation

Leetcode Problem:

• Design a Stack With Increment Operation

Summary

• Design a stack that supports increment operations on its elements
• The CustomStack class is implemented with a maximum size and supports push, pop, and increment operations.

Approach

• The CustomStack class uses an array to store the stack elements and a separate array to store the lazy values
• The lazy array is used to store the increment values for each element in the stack
• The push operation adds an element to the top of the stack, the pop operation removes an element from the top of the stack and updates the lazy values, and the increment operation updates the lazy values for a range of elements in the stack.

Complexity

• O(1) for push, O(1) for pop, and O(min(k, top)) for increment operations

Explain

• The CustomStack class is implemented with two arrays, arr and lazy, to store the stack elements and lazy values respectively
• The top variable is used to keep track of the top element in the stack
• The push operation adds an element to the top of the stack and updates the lazy value for the top element
• The pop operation removes an element from the top of the stack and updates the lazy values for the elements below it
• The increment operation updates the lazy values for a range of elements in the stack
• The time complexity of the push and pop operations is O(1) because they only involve updating a single element in the stack
• The time complexity of the increment operation is O(min(k, top)) because it involves updating the lazy values for a range of elements in the stack.

Solution Code:

class CustomStack {
public:
    int stackSize;
    int arr[1001];
    int lazy[1001];
    int top;
    CustomStack(int maxSize) {
        stackSize = maxSize;
        top = 0;
        for(int i = 0 ; i < maxSize; i++){
            lazy[i] = 0;
        }
    }
    
    void push(int x) {
        if (top == stackSize) return;
        arr[top++] = x;
    }
    
    int pop() {
        if (top == 0) return -1;
        top--;
        if (top > 0) lazy[top-1] += lazy[top];
        arr[top] += lazy[top];
        lazy[top] = 0;
        return arr[top];
    }
    
    void increment(int k, int val) {
        if (top == 0) return;
        lazy[min(k-1, top-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);
 */