Design a Stack With Increment Operation

Design a Stack With Increment Operation

Problem Description
Design a stack that supports, in addition to the standard push, pop, and top operations, an increment operation. The increment operation takes two parameters: k and val, which means that the bottom k elements of the stack (if the stack has fewer than k elements, apply to all elements) should each be incremented by val. All operations should have a time complexity of O(1).

Solution Approach
Standard stack operations are O(1) themselves, but if the increment operation directly traverses the bottom k elements, the time complexity would be O(k), which does not meet the requirement. To achieve O(1) increment operations, we need to adopt a lazy update approach—record the increment operations and apply them only when elements are popped.

Detailed Steps

  1. Data Structure Design

    • Use an array (or list) stack to store the stack elements.
    • Use another array (or list) inc to record increment operations. inc[i] represents the total increment that should be applied to all elements from the bottom of the stack up to position i (inclusive).
    • Maintain a variable size to indicate the current number of elements in the stack.

  2. Push Operation

    • Push the new element into stack.
    • Initialize the corresponding position in the inc array (i.e., position size) to 0, as the new element has not yet been affected by any increment operations.
    • Increase size.

  3. Pop Operation

    • If the stack is empty, return an error.
    • Otherwise, first calculate the actual value of the current top element. Since inc[size-1] records the total increment for all elements from the bottom to the top element, the actual value is stack[size-1] + inc[size-1].
    • Note: This increment value only applies to the current element, but the elements below the top may have pending increments. Therefore, before popping the top element, we need to add inc[size-1] to inc[size-2] (if it exists). This ensures that when the element below the top is popped later, it will inherit this increment.
    • Finally, pop the top elements from both stack and inc, and decrease size.

  4. Top Operation

    • Similar to pop, but without removing the element. Directly return stack[size-1] + inc[size-1].

  5. Increment Operation

    • Determine the actual number of elements to increment: m = min(k, size).
    • Add the increment val to inc[m-1]. Because inc[m-1] records the total increment for all elements from the bottom to the m-th element, when these elements are popped, they will correctly apply the increment through the propagation mechanism in the pop operation.

Example
Assume the following operations are executed in order:

  1. push(1)
  2. push(2)
  3. increment(2, 100) // Add 100 to the bottom 2 elements
  4. push(3)
  5. increment(2, 50) // Add 50 to the bottom 2 elements (now the stack has 3 elements, the bottom two are 1 and 2)
  6. pop() // Pop 3
  7. pop() // Pop 2
  8. pop() // Pop 1

Execution process:

  • After operation 1: stack = [1], inc = [0], size=1
  • After operation 2: stack = [1,2], inc = [0,0], size=2
  • Operation 3: k=2, val=100 -> inc[1] += 100 → inc = [0,100]
  • After operation 4: stack = [1,2,3], inc = [0,100,0], size=3
  • Operation 5: k=2, val=50 -> inc[1] += 50 → inc = [0,150,0]
  • Operation 6 (pop):
    • Top index=2, actual value=3+0=3
    • inc[1] += inc[2] → inc = [0,150,0] (no change because inc[2]=0)
    • After pop: stack = [1,2], inc = [0,150], size=2
  • Operation 7 (pop):
    • Top index=1, actual value=2+150=152
    • inc[0] += inc[1] → inc = [150,150]
    • After pop: stack = [1], inc = [150], size=1
  • Operation 8 (pop):
    • Top index=0, actual value=1+150=151
    • After pop, the stack is empty

Final pop order: 3, 152, 151, as expected (1 and 2 were each incremented by 100+50=150).