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.

Example Walkthrough

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

1. Initialization

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

2. Push Operations

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

3. Increment Operation

   1increment(2, 100) -> stk = [5, 7, 3], add = [100, 100, 0], i = 3
   2(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

   1pop() -> stk = [5, 7, 0], add = [100, 100, 0], i = 2
   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.)
   3
   4pop() -> stk = [5, 0, 0], add = [100, 0, 0], i = 1
   5(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.)
   6
   7pop() -> stk = [0, 0, 0], add = [0, 0, 0], i = 0
   8(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

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.