1381. Design a Stack With Increment Operation


Problem Description

This problem requires designing a stack that not only has the usual push and pop functionalities but also a special feature that allows incrementing the bottom-most k elements by a given value. The class CustomStack should be initialized with a maxSize, which is the maximum number of elements that the stack can hold. The stack supports three operations:

  1. push(int x): This method should push the value x onto the stack only if the number of elements in the stack is less than the maxSize.

  2. pop(): This method should remove the top element of the stack and return its value. If the stack is empty, it should return -1.

  3. increment(int k, int val): This method should increment the bottom k elements of the stack by the value val. If there are fewer than k elements in the stack, it increments all of them.

Intuition

To solve this problem efficiently, one could keep a secondary array, called add, alongside the primary stack data structure, which is an array called stk. The array add should have the same length as stk, and it's used to store the incremental values.

The key idea behind the solution is to use lazy incrementation. When the increment method is called, instead of incrementing every single one of the bottom k elements, the value is added to the add array at the index that corresponds to the k-th element (or the last element if k is greater than the number of elements in the stack). When an element is popped, the increment value at the current index is added to the popped element to reflect all the increments up to that point. If there are elements below it, the increment is passed down to the next element in the stack, ensuring all the increments are applied lazily during the pop operations.

This approach is efficient because it does not require incrementing each element individually, which could be costly for a large number of inc operations or a large k. Instead, the increment is deferred until an element is popped, making both push and increment operations O(1) time complexity, with the pop operation being O(1) as well.

Learn more about Stack patterns.

Solution Approach

The implementation of CustomStack relies on two arrays, stk to actually hold the stack elements, and add to keep track of the increment values that need to be applied to each element. The helper variable i acts as a pointer to the current top of the stack, with i = 0 indicating that the stack is empty.

Here's how each method of the CustomStack works according to the solution approach:

__init__(self, maxSize: int)

The constructor initializes the two arrays, stk and add, with a length of maxSize and fills them with zeros. It also initializes i to 0, which represents the initial position of the stack pointer.

push(self, x: int) -> None

The push method checks whether the current stack size (i) is less than the maximum permissible size before proceeding to add an element. If there's space, the method places the value x at the current top position (given by i) in stk and increments i by 1 to reflect the new top of the stack.

pop(self) -> int

The pop method first checks if the stack is empty by checking if i is 0 or less. If the stack is empty, it returns -1. If not, it decrements i by 1 to point to the current top of the stack, retrieves the original value plus any increment value stored in add for that position, and returns it. If there is another element below the popped element (i > 0), the increment value for the current position is passed down to the next element by adding it to the add value of the next (now the top) element. Finally, the increment value of the popped position is reset to 0.

increment(self, k: int, val: int) -> None

The increment operation is done lazily. Instead of iterating through the bottom k elements, it finds the position i in the add array which corresponds to the k-th element or the last element if the stack's current size is less than k. It then adds the val to this add[i]. When a pop operation is executed later, this value will be added to the returned element and carried over to the next element if necessary.

The data structures used, namely the two arrays stk and add, are efficient for the operations that the CustomStack class is expected to perform. The lazy increment pattern avoids unnecessary work during increment operations, ensuring that the pop and push operations remain efficient even with a large number of increments.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's illustrate the solution approach with a small example. Assume we initialize a CustomStack with maxSize = 3.

  1. Initialization

    CustomStack(maxSize = 3)
    stk = [0, 0, 0]
    add = [0, 0, 0]
    i   = 0 (stack is empty at this stage)
  2. Push Operations

    push(5) -> stk = [5, 0, 0], add = [0, 0, 0], i = 1
    push(7) -> stk = [5, 7, 0], add = [0, 0, 0], i = 2
    push(3) -> stk = [5, 7, 3], add = [0, 0, 0], i = 3
  3. Increment Operation

    increment(2, 100) -> stk = [5, 7, 3], add = [100, 100, 0], i = 3
    (We added 100 to the first two elements, however, we store this increment in the `add` array rather than change `stk` directly.)
  4. Pop Operation

    pop() -> stk = [5, 7, 0], add = [100, 100, 0], i = 2
    (We pop the top element, which would be 3 + 100 from `add[2]` = 103, and since there's no element below it, we don't need to transfer any increment value.)
    
    pop() -> stk = [5, 0, 0], add = [100, 0, 0], i = 1
    (We pop off 7 + 100 from `add[1]` = 107, and we pass the increment value 100 to the next element by setting `add[0]` to 200.)
    
    pop() -> stk = [0, 0, 0], add = [0, 0, 0], i = 0
    (The final pop is 5 + 200 from `add[0]` = 205, and we set `add[0]` back to 0 as it’s no longer needed.)

Based on this example, we can observe that the increment operation updates the add array rather than every element, and the incremented amounts are only applied when pop operations occur. This makes accurate use of the lazy incrementation strategy where the increment operation is O(1) and doesn't depend on the number of elements it affects. The pop operation remains O(1) because it only involves a constant number of actions regardless of the size of the stack or the increment values.

Solution Implementation

1class CustomStack:
2    def __init__(self, max_size: int):
3        # Initialize the stack with the given max size.
4        # The actual stack is maintained in a list initialized with zeros.
5        self.stack = [0] * max_size
6        # The add list is used to store the additive increments for each position.
7        self.add = [0] * max_size
8        # The index `current_size` tracks the number of elements in the stack.
9        self.current_size = 0
10
11    def push(self, x: int) -> None:
12        # Push an element onto the stack if there is space available.
13        if self.current_size < len(self.stack):
14            self.stack[self.current_size] = x
15            self.current_size += 1
16
17    def pop(self) -> int:
18        # Pop the top element from the stack and apply any increments to it.
19        if self.current_size <= 0:
20            return -1  # Return -1 if the stack is empty.
21        self.current_size -= 1
22        result = self.stack[self.current_size] + self.add[self.current_size]
23        # Transfer the increment to the next element to be popped, if applicable.
24        if self.current_size > 0:
25            self.add[self.current_size - 1] += self.add[self.current_size]
26        # Reset the increment for the current position.
27        self.add[self.current_size] = 0
28        return result  # Return the final value after applying the increment.
29
30    def increment(self, k: int, val: int) -> None:
31        # Increment the bottom `k` elements of the stack by `val`.
32        limit = min(k, self.current_size) - 1  # Determine the actual limit to increment.
33        if limit >= 0:
34            self.add[limit] += val  # Apply the increment to the `limit` position.
35
36# Example usage:
37# max_size = 3
38# custom_stack = CustomStack(max_size)
39# custom_stack.push(1)
40# custom_stack.push(2)
41# print(custom_stack.pop())  # Outputs: 2
42# custom_stack.push(2)
43# custom_stack.increment(2, 1)  # Stack becomes [2,3]
44# print(custom_stack.pop())  # Outputs: 3
45# print(custom_stack.pop())  # Outputs: 2
46
1class CustomStack {
2    private int[] stack;  // Array to store stack elements
3    private int[] increments;  // Array to store increment operations
4    private int topIndex;  // Points to the next free spot in the 'stack' array (and the current size of the stack)
5
6    // Constructor to initialize the stack and increments array with the given maxSize
7    public CustomStack(int maxSize) {
8        stack = new int[maxSize];
9        increments = new int[maxSize];
10        topIndex = 0;
11    }
12
13    // Method to push an element onto the top of the stack if there is space available
14    public void push(int x) {
15        if (topIndex < stack.length) {
16            stack[topIndex++] = x;
17        }
18    }
19
20    // Method to remove and return the top element of the stack; if the stack is empty, returns -1
21    public int pop() {
22        if (topIndex <= 0) { // Check if stack is empty
23            return -1;
24        }
25        int result = stack[--topIndex] + increments[topIndex]; // Get the top element with the increment
26        if (topIndex > 0) { // If there's still an element below the top, transfer the increment
27            increments[topIndex - 1] += increments[topIndex];
28        }
29        increments[topIndex] = 0; // Reset the increment for the current index since it's been applied
30        return result;
31    }
32
33    // Method to increment the bottom 'k' elements by 'val'
34    public void increment(int k, int val) {
35        if (topIndex > 0) { // Check if stack is not empty
36            int index = Math.min(topIndex, k) - 1; // Determine the index until which the increments should be applied
37            increments[index] += val; // Apply the increment to the kth element or the current top, whichever is smaller
38        }
39    }
40}
41
42/**
43 * The following operations can be performed on an instance of CustomStack:
44 * - Create a new stack with a maximum size `maxSize`.
45 * - Push an integer `x` onto the stack if there is space.
46 * - Pop and return the top element of the stack; if the stack is empty, return -1.
47 * - Increment the bottom `k` elements of the stack by `val`.
48 *
49 * Usage:
50 * CustomStack customStack = new CustomStack(maxSize);
51 * customStack.push(x);
52 * int topElement = customStack.pop();
53 * customStack.increment(k, val);
54 */
55
1#include <vector>
2#include <algorithm> // For min function
3
4class CustomStack {
5public:
6    // Constructor to initialize the stack with a maximum size
7    CustomStack(int maxSize) {
8        stack_.resize(maxSize);
9        additional_.resize(maxSize);
10        size_ = 0; // Stack starts empty
11    }
12
13    // Method to push an element to the top of the stack if there's space
14    void push(int x) {
15        if (size_ < stack_.size()) {
16            stack_[size_] = x; // Insert element at the top of the stack
17            size_++; // Increase the stack size
18        }
19    }
20
21    // Method to pop the top element from the stack
22    int pop() {
23        // Check if the stack is empty
24        if (size_ <= 0) {
25            return -1; // If empty, return -1 as per the problem statement
26        }
27        // Decrement size and get the top value with additional increment applied
28        size_--;
29        int value = stack_[size_] + additional_[size_];
30        if (size_ > 0) {
31            // Propagate the increment to the below element (if any)
32            additional_[size_ - 1] += additional_[size_];
33        }
34        // Reset the additional increment for the popped element
35        additional_[size_] = 0;
36        return value; // Return the popped value
37    }
38
39    // Method to increment the bottom k elements of the stack by val
40    void increment(int k, int val) {
41        if (size_ > 0) {
42            // Only apply to the bottom k elements or all elements if size is less than k
43            int limit = std::min(k, size_);
44            additional_[limit - 1] += val; // Apply increment to the k-th element from the bottom
45        }
46    }
47
48private:
49    std::vector<int> stack_;     // Vector to store the stack elements
50    std::vector<int> additional_; // Vector to store additional increments to apply on pop
51    int size_;                   // Current size of the stack (number of elements)
52};
53
54// Below this line is where the CustomStack class could be used:
55/*
56int main() {
57    CustomStack* obj = new CustomStack(maxSize);
58    obj->push(x);
59    int param_2 = obj->pop();
60    obj->increment(k,val);
61  
62    // Remember to delete the object created with 'new' to prevent memory leaks
63    delete obj;
64    return 0;
65}
66*/
67
1// Global variables to act as the stack, the "add" track array, and the index pointer.
2let stack: number[];
3let addTrack: number[];
4let currentIndex: number;
5
6// Function to initialize the stack with a maximum size.
7function initializeCustomStack(maxSize: number): void {
8    stack = Array(maxSize).fill(0);
9    addTrack = Array(maxSize).fill(0);
10    currentIndex = 0;
11}
12
13// Function to push a new element onto the stack if there's space.
14function pushToCustomStack(value: number): void {
15    if (currentIndex < stack.length) {
16        stack[currentIndex++] = value;
17    }
18}
19
20// Function to pop the top element from the stack, with additional logic for increment overflow.
21function popFromCustomStack(): number {
22    if (currentIndex <= 0) {
23        return -1;
24    }
25    const topValue = stack[--currentIndex] + addTrack[currentIndex];
26    if (currentIndex > 0) {
27        addTrack[currentIndex - 1] += addTrack[currentIndex];
28    }
29    addTrack[currentIndex] = 0;
30    return topValue;
31}
32
33// Function to increment the bottom k elements of the stack by a specified value.
34function incrementBottomElements(k: number, value: number): void {
35    let incrementIndex = Math.min(currentIndex, k) - 1;
36    if (incrementIndex >= 0) {
37        addTrack[incrementIndex] += value;
38    }
39}
40
41// Please note that these global functions assume that the initializeCustomStack function is called first to set up the stack.
42

Time and Space Complexity

Time Complexity:

  • init: O(n) where n is maxSize. This is because we're initializing two arrays stk and add, each of maxSize length, with zeroes.

  • push: O(1) for each operation. Inserting an element into the stack is done by directly indexing into the array, which is a constant-time operation.

  • pop: O(1) for each operation. Removing an element from the stack is also done by directly indexing into the array. The additional operation of transferring the increment value to the next element is constant-time as well, as it involves only a direct index and an arithmetic operation.

  • increment: O(1) for each operation. Even though the increment operation affects k elements, we're lazily applying this operation by only updating one item in the add array, so we're not iterating through the entire k elements at the time of calling this method.

Space Complexity:

  • Overall space complexity is O(n) where n is maxSize. This is due to maintaining two additional arrays (stk and add) that store the elements of the stack and the increment operations to be applied.

Learn more about how to find time and space complexity quickly using problem constraints.


Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:
Question 1 out of 10

Which algorithm should you use to find a node that is close to the root of the tree?


Recommended Readings

Want a Structured Path to Master System Design Too? Don’t Miss This!


Load More