Problem Description

You need to design a custom stack data structure that supports standard stack operations plus an increment operation.

The CustomStack class should support:

1. Constructor CustomStack(int maxSize): Creates a stack with a maximum capacity of maxSize elements. The stack cannot hold more than maxSize elements.

2. Push operation push(int x): Adds element x to the top of the stack. If the stack has already reached its maximum capacity (maxSize), the element should not be added.

3. Pop operation pop(): Removes and returns the element from the top of the stack. If the stack is empty, returns -1.

4. Increment operation increment(int k, int val): Increases the value of the bottom k elements in the stack by val. If the stack contains fewer than k elements, all elements in the stack should be incremented by val.

For example, if the stack contains [1, 2, 3] (with 3 at the top) and you call increment(2, 100), the stack becomes [101, 102, 3] because the bottom 2 elements are incremented by 100.

The solution uses an optimization technique with a lazy propagation approach. Instead of directly modifying all elements during the increment operation, it uses an additional array add to track cumulative increments at each position. This allows the increment operation to run in O(1) time. When popping an element, the accumulated increment value is applied and propagated to the element below it if necessary.

Intuition

The straightforward approach would be to use a simple array to simulate the stack and directly modify elements during the increment operation. However, this would make the increment operation O(k) in time complexity, as we'd need to iterate through and update k elements.

The key insight is that we can defer the actual increments until we need to access the elements. Think about it this way: if we increment the bottom k elements by some value, and then increment them again before popping, we're doing redundant work by updating the same elements multiple times.

Instead of immediately applying the increment to all k elements, we can just "mark" that an increment needs to be applied. We use a separate array add where add[i] stores the cumulative increment that should be applied to the element at position i when it's accessed.

The clever part comes with how we handle these increments during the pop operation. When we pop an element at position i, we:

1. Apply the increment stored at add[i] to get the actual value
2. Pass this increment down to the element below it (add[i-1] += add[i])
3. Clear the increment at position i since we've already used it

This "passing down" mechanism ensures that any increments meant for elements below the popped element are preserved. For example, if we increment the bottom 3 elements and then pop the 3rd element, the remaining 2 elements still need to remember they should be incremented.

By using this lazy propagation technique, we achieve O(1) time complexity for all operations: push and pop work with single elements, and increment only updates a single position in the add array rather than iterating through k elements.

Learn more about Stack patterns.

Solution Approach

The implementation uses two arrays and an index pointer to efficiently manage the custom stack with lazy increment propagation.

Data Structures:

• stk: An array of size maxSize to store the actual stack elements
• add: An array of size maxSize to store cumulative increment values at each position
• i: An integer pointer indicating the next position to insert an element (also represents the current size of the stack)

Implementation Details:

1. Initialization (__init__):

   • Create stk array with maxSize capacity, initialized with zeros
   • Create add array with maxSize capacity, initialized with zeros
   • Set i = 0 to indicate the stack is empty

2. Push Operation (push):

   • Check if there's space: if self.i < len(self.stk)
   • If yes, place element x at position stk[i]
   • Increment i by 1 to update the stack size
   • Time complexity: O(1)

3. Pop Operation (pop):

   • Check if stack is empty: if self.i <= 0, return -1
   • Decrement i by 1 to point to the top element
   • Calculate the actual value: ans = stk[i] + add[i]
   • Propagate the increment downward: if there's an element below (i > 0), add the current increment to it: add[i-1] += add[i]
   • Clear the current position's increment: add[i] = 0
   • Return the calculated answer
   • Time complexity: O(1)

4. Increment Operation (increment):

   • Calculate the actual position to increment: i = min(k, self.i) - 1
   • This ensures we don't exceed the current stack size
   • If i >= 0 (stack is not empty), add val to add[i]
   • This marks that all elements from bottom to position i should be incremented by val
   • Time complexity: O(1)

Key Pattern - Lazy Propagation: The solution employs lazy propagation to defer the actual increment operations. Instead of immediately updating all k elements, we only mark the topmost affected position with the increment value. When elements are popped, the increment values cascade down to the remaining elements, ensuring correctness while maintaining O(1) complexity for all operations.

For example, if we have stack [1, 2, 3] and call increment(2, 100), instead of immediately updating to [101, 102, 3], we just set add[1] = 100. When we pop elements, the actual values are computed on-the-fly by adding the base value and the accumulated increment.

Example Walkthrough

Let's walk through a sequence of operations on a CustomStack with maxSize = 3:

Initial State:

• stk = [0, 0, 0]
• add = [0, 0, 0]
• i = 0 (stack is empty)

Operation 1: push(5)

• Place 5 at position 0: stk = [5, 0, 0]
• Increment i: i = 1
• State: stk = [5, 0, 0], add = [0, 0, 0]

Operation 2: push(7)

• Place 7 at position 1: stk = [5, 7, 0]
• Increment i: i = 2
• State: stk = [5, 7, 0], add = [0, 0, 0]

Operation 3: push(10)

• Place 10 at position 2: stk = [5, 7, 10]
• Increment i: i = 3
• State: stk = [5, 7, 10], add = [0, 0, 0]

Operation 4: increment(2, 100)

• Want to increment bottom 2 elements
• Calculate position: min(2, 3) - 1 = 1
• Set add[1] = 100 (marking that bottom 2 elements need +100)
• State: stk = [5, 7, 10], add = [0, 100, 0]
• Note: Actual values are [5, 107, 10] but we haven't modified stk yet

Operation 5: pop()

• Decrement i: i = 2 (pointing to top element)
• Calculate value: stk[2] + add[2] = 10 + 0 = 10
• No increment to propagate since add[2] = 0
• Clear: add[2] = 0 (already 0)
• Return: 10
• State: stk = [5, 7, 10], add = [0, 100, 0], i = 2

Operation 6: pop()

• Decrement i: i = 1
• Calculate value: stk[1] + add[1] = 7 + 100 = 107
• Propagate increment down: add[0] = add[0] + add[1] = 0 + 100 = 100
• Clear: add[1] = 0
• Return: 107
• State: stk = [5, 7, 10], add = [100, 0, 0], i = 1

Operation 7: increment(1, 50)

• Want to increment bottom 1 element
• Calculate position: min(1, 1) - 1 = 0
• Add to existing increment: add[0] = 100 + 50 = 150
• State: stk = [5, 7, 10], add = [150, 0, 0], i = 1

Operation 8: pop()

• Decrement i: i = 0
• Calculate value: stk[0] + add[0] = 5 + 150 = 155
• No element below to propagate to (i = 0)
• Clear: add[0] = 0
• Return: 155
• State: stk = [5, 7, 10], add = [0, 0, 0], i = 0 (empty stack)

This example demonstrates how the lazy propagation works: increments are stored in the add array and only applied when elements are popped, with the increment values cascading down to lower elements as needed.

Solution Implementation

Time and Space Complexity

Time Complexity: O(1) for all operations (push, pop, and increment)

• push: Simply assigns a value to self.stk[self.i] and increments the index, which takes constant time.
• pop: Performs a fixed number of operations - decrements the index, calculates the result by adding self.stk[self.i] + self.add[self.i], potentially updates self.add[self.i - 1], and resets self.add[self.i]. All these operations take constant time.
• increment: Calculates min(k, self.i) - 1 and updates a single element in the self.add array. This is a constant time operation regardless of the value of k.

The key insight is that the increment operation uses lazy propagation. Instead of updating all k elements immediately, it only marks the increment at position i in the self.add array. The actual increment is applied when elements are popped, and the increment value is propagated down the stack one level at a time during each pop operation.

Space Complexity: O(n) where n is the maxSize parameter

• The implementation uses two arrays: self.stk and self.add, each of size maxSize.
• The space usage is 2 * maxSize = O(maxSize) = O(n).
• Additionally, there's a constant amount of space for the index variable self.i.

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

Common Pitfalls

1. Forgetting to Propagate Increments During Pop

One of the most common mistakes is forgetting to propagate the increment value to the element below when popping. This breaks the lazy propagation mechanism.

Incorrect Implementation:

def pop(self) -> int:
    if self.current_size <= 0:
        return -1
  
    self.current_size -= 1
    result = self.stack[self.current_size] + self.increment_values[self.current_size]
  
    # MISTAKE: Forgetting to propagate increment to element below
    self.increment_values[self.current_size] = 0  # Only clearing without propagation
  
    return result

Correct Implementation:

def pop(self) -> int:
    if self.current_size <= 0:
        return -1
  
    self.current_size -= 1
    result = self.stack[self.current_size] + self.increment_values[self.current_size]
  
    # Propagate increment to the element below before clearing
    if self.current_size > 0:
        self.increment_values[self.current_size - 1] += self.increment_values[self.current_size]
  
    self.increment_values[self.current_size] = 0
  
    return result

2. Off-by-One Error in Increment Operation

Another frequent mistake is incorrectly calculating the index for the increment operation, especially when dealing with 1-based k versus 0-based indexing.

Incorrect Implementation:

def increment(self, k: int, val: int) -> None:
    # MISTAKE: Using k directly as index without adjustment
    target_index = min(k, self.current_size)
  
    if target_index >= 0:
        self.increment_values[target_index] += val  # This would increment k+1 elements!

Correct Implementation:

def increment(self, k: int, val: int) -> None:
    # Correctly subtract 1 to convert from count to index
    target_index = min(k, self.current_size) - 1
  
    if target_index >= 0:
        self.increment_values[target_index] += val

3. Not Clearing Increment Values After Pop

Failing to reset the increment value at a popped position can cause incorrect values if that position is reused later.

Incorrect Implementation:

def pop(self) -> int:
    if self.current_size <= 0:
        return -1
  
    self.current_size -= 1
    result = self.stack[self.current_size] + self.increment_values[self.current_size]
  
    if self.current_size > 0:
        self.increment_values[self.current_size - 1] += self.increment_values[self.current_size]
  
    # MISTAKE: Not clearing the increment value
    # self.increment_values[self.current_size] = 0  # Missing this line!
  
    return result

This would cause issues when a new element is pushed to the same position later, as it would inherit the old increment value.

4. Incorrect Boundary Check in Increment

Not properly handling the case when the stack is empty or when k is 0.

Incorrect Implementation:

def increment(self, k: int, val: int) -> None:
    # MISTAKE: Not checking if the result is negative
    target_index = min(k, self.current_size) - 1
    self.increment_values[target_index] += val  # Could access index -1!

Correct Implementation:

def increment(self, k: int, val: int) -> None:
    target_index = min(k, self.current_size) - 1
  
    # Proper boundary check
    if target_index >= 0:
        self.increment_values[target_index] += val

Prevention Tips:

1. Always trace through the lazy propagation logic with small examples
2. Verify that increment values are properly propagated and cleared
3. Test edge cases: empty stack, single element, k larger than stack size
4. Use consistent naming and clear comments to track the propagation flow
5. Remember that the increment array represents cumulative increments that apply to all elements from bottom up to that index