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:
-
push(int x)
: This method should push the valuex
onto the stack only if the number of elements in the stack is less than themaxSize
. -
pop()
: This method should remove the top element of the stack and return its value. If the stack is empty, it should return-1
. -
increment(int k, int val)
: This method should increment the bottomk
elements of the stack by the valueval
. If there are fewer thank
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 EvaluatorExample Walkthrough
Let's illustrate the solution approach with a small example. Assume we initialize a CustomStack
with maxSize = 3
.
-
Initialization
CustomStack(maxSize = 3) stk = [0, 0, 0] add = [0, 0, 0] i = 0 (stack is empty at this stage)
-
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
-
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.)
-
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)
wheren
ismaxSize
. This is because we're initializing two arraysstk
andadd
, each ofmaxSize
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 affectsk
elements, we're lazily applying this operation by only updating one item in theadd
array, so we're not iterating through the entirek
elements at the time of calling this method.
Space Complexity:
- Overall space complexity is
O(n)
wheren
ismaxSize
. This is due to maintaining two additional arrays (stk
andadd
) 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.
Which algorithm should you use to find a node that is close to the root of the tree?
Recommended Readings
Stack Intro Imagine you have a pile of books on your desk If you want to add a new book you place it on top If you want to read a book you take it from the top And if you simply want to see which book is on the top you
LeetCode Patterns Your Personal Dijkstra's Algorithm to Landing Your Dream Job The goal of AlgoMonster is to help you get a job in the shortest amount of time possible in a data driven way We compiled datasets of tech interview problems and broke them down by patterns This way we
Recursion Recursion is one of the most important concepts in computer science Simply speaking recursion is the process of a function calling itself Using a real life analogy imagine a scenario where you invite your friends to lunch https algomonster s3 us east 2 amazonaws com recursion jpg You first
Want a Structured Path to Master System Design Too? Don’t Miss This!