Design A Stack With Increment Operation - Leetcode Solution

Problem Statement:

Design a stack that supports an increment operation, which adds a given integer increment to the bottom k elements of the stack and returns the current top element.

Implement the CustomStack class:

CustomStack(int maxSize) Initializes the object with maxSize which is the maximum number of elements in the stack or do nothing if the stack reached the maximum capacity. void push(int x) Adds x to the top of the stack if the stack hasn't 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, just increment all the elements in the stack.

Solution:

We will begin by creating a class called CustomStack which will contain three functions: push, pop, and inc. We will also create a vector called stack to store the elements of the stack. The size of the stack will be initialized using the maxSize parameter.

class CustomStack { private: vector<int> stack; int maxSize;

public: CustomStack(int maxSize) { this->maxSize = maxSize; }

void push(int x) {
    if(stack.size() < maxSize) {
        stack.push_back(x);
    }
}

int pop() {
    if(stack.size() > 0) {
        int top_element = stack.back();
        stack.pop_back();
        return top_element;
    } else {
        return -1;
    }
}

void inc(int k, int val) {
    
}

};

Now, let's implement the inc function which will add the given integer increment to the bottom k elements of the stack and return the current top element. We will iterate through the first k elements of the stack and add the given integer increment to each element. If there are less than k elements in the stack, we will add the given integer increment to all the elements in the stack. We will also use a vector called increment which will store the increments for each index in the stack. This vector will be used later to update the bottom elements of the stack.

void inc(int k, int val) { if(stack.size() < k) { k = stack.size(); } for(int i = 0; i < k; i++) { stack[i] += val; } increment.push_back(val); }

Now, let's modify the push and pop functions to reflect the changes made in the inc function. When we pop an element off the stack, we will also subtract the corresponding element in the increment vector. When we push an element onto the stack, we will push the sum of the increment for that index and the given value onto the stack. If the size of the stack is greater than or equal to the maxSize, we will not add any element to the stack.

void push(int x) { if(stack.size() < maxSize) { int index = stack.size(); if(index < increment.size()) { x += increment[index]; } stack.push_back(x); } }

int pop() { if(stack.size() > 0) { int top_element = stack.back(); if(stack.size()-1 < increment.size()) { increment[stack.size()-1] -= increment[stack.size()-1]; } stack.pop_back(); return top_element; } else { return -1; } }

The final program implementing the CustomStack class with all functions can be seen here:

class CustomStack { private: vector<int> stack; vector<int> increment; int maxSize;

public: CustomStack(int maxSize) { this->maxSize = maxSize; }

void push(int x) {
    if(stack.size() < maxSize) {
        int index = stack.size();
        if(index < increment.size()) {
            x += increment[index];
        }
        stack.push_back(x);
    }
}

int pop() {
    if(stack.size() > 0) {
        int top_element = stack.back();
        if(stack.size()-1 < increment.size()) {
            increment[stack.size()-1] -= increment[stack.size()-1];
        }
        stack.pop_back();
        return top_element;
    } else {
        return -1;
    }
}

void inc(int k, int val) {
    if(stack.size() < k) {
        k = stack.size();
    }
    for(int i = 0; i < k; i++) {
        stack[i] += val;
    }
    increment.push_back(val);
}

};

Time Complexity:

The push and pop functions have time complexity O(1). The inc function has time complexity O(k) as it needs to iterate through the first k elements of the stack. Overall, the time complexity for this program is O(k*n) where n is the maximum size of the stack and k is the maximum number of elements to increment.

1