Design a Stack That Supports Increment Operations

Design a Stack That Supports Increment Operations

Problem Description
Design a stack data structure that, in addition to supporting conventional stack operations (push, pop, top), also supports an increment operation increment(k, val). This operation adds val to the values of the bottom k elements of the stack. If the number of elements in the stack is less than k, then all elements are increased by val. All operations are required to have O(1) time complexity.

Solution Process

1. Problem Analysis
Conventional stack operations push, pop, and top all have O(1) time complexity. However, if the increment(k, val) operation is implemented directly, it would require modifying the bottom k elements, which in the worst case takes O(n) time (when k is close to n), failing to meet the problem's requirements. Therefore, we need a lazy update mechanism that defers the effect of the increment operation until the element is accessed (via pop or top).

2. Core Idea: Lazy Update
We can use two auxiliary stacks:

  • A main stack stack to store the actual values.
  • An increment stack inc to record deferred increment operations. inc[i] represents the cumulative increment value that needs to be applied to all elements from the bottom of the stack up to position i.

A more efficient approach is to record only the differences of the increment operations. Specifically:

  • We use an array or linked list as the structure for the stack.
  • We additionally maintain an array add, where add[i] represents the increment value for the segment from the bottom of the stack to the i-th element (including i).
  • When executing increment(k, val), we only need to update two positions in the add array (actually, indices in the stack) instead of modifying each element in the stack.

3. Detailed Design
We use an array data to simulate the stack and maintain a stack top pointer top (pointing to the next available insertion position). Simultaneously, we maintain an equally long array add to record increments.

The key points are:

  • When executing increment(k, val):
    • Let idx = min(k, top) - 1, i.e., the deepest affected position (bottom index is 0, top index is top-1).
    • Increase add[idx] by val. This way, when we later access an element, we can obtain its final value by accumulating the increments from that position to the top of the stack.

However, when popping, we need to know the sum of all increments from the current element to the top of the stack. If we sum from the beginning each time, the worst-case scenario is O(n). To maintain O(1), we need a different strategy: Difference Array.

4. Difference Array Optimization
We use a difference array diff, where diff[i] represents the difference between add[i] and add[i-1] (add[-1] is considered 0). Thus:

  • increment(k, val) operation:
    • Let idx = min(k, top) - 1.
    • Increase diff[idx] by val.
  • When we need to know the cumulative increment at position i, add[i] = diff[0] + diff[1] + ... + diff[i]. This again becomes O(n) summation.

To obtain the cumulative increment at any position in O(1) time, we can dynamically maintain a cumulative variable sum during pop, accumulating diff from the top to the bottom. However, this makes the pop operation not O(1). We need a more clever approach.

5. Final Solution: Lazy Propagation
We maintain an auxiliary stack inc of the same size as the main stack. inc[i] represents the increment value starting from position i towards the bottom of the stack. More precisely:

  • When executing increment(k, val):
    • Let idx = min(k, top) - 1.
    • If idx >= 0, then increase inc[idx] by val.
  • When executing pop():
    • Pop the top element data[top-1] and the corresponding inc[top-1].
    • If inc[top-1] is not 0, we "propagate" this increment value to the next element (i.e., the new top element top-2): inc[top-2] += inc[top-1].
    • Then return data[top-1] + cumulative increment (note that due to the propagation mechanism, the popped element already carries all the values it should have added).

Actually, we can simplify: use only one inc array and perform propagation during pop. The top operation also requires similar handling, but care must be taken not to modify the stack state.

6. Specific Implementation Steps
We use a dynamic array (such as a vector) to store data and a pointer top to represent the top position. Simultaneously, we maintain an equally long inc array.

  • Initialization: Create empty arrays data and inc, top = 0.
  • push(val): Place val into data[top], set inc[top] = 0, then top++.
  • pop():
    1. If top == 0, the stack is empty; report an error.
    2. top--.
    3. If top > 0, propagate the increment of the current top element to the new top element: inc[top-1] += inc[top].
    4. Return data[top] + inc[top] (note that top now points to the position of the popped element).
  • top():
    1. If top == 0, the stack is empty; report an error.
    2. Return data[top-1] + inc[top-1].
  • increment(k, val):
    1. If top == 0, return directly (no elements to increment).
    2. Let idx = min(k, top) - 1.
    3. Increase inc[idx] by val.

7. Time Complexity Analysis

  • push: O(1)
  • pop: O(1) (although there is an addition operation, it's constant time)
  • top: O(1)
  • increment: O(1)

All operations satisfy O(1) time complexity.

8. Example Demonstration
Assume the following sequence of operations:

  1. push(1); push(2); push(3);
    data = [1,2,3], inc = [0,0,0], top=3
  2. increment(2, 100);
    k=2, idx=min(2,3)-1=1, inc[1] +=100 → inc=[0,100,0]
  3. top() → returns data[2]+inc[2]=3+0=3
  4. pop() →
    top-- → top=2
    inc[1] += inc[2] → inc[1]=100+0=100
    returns data[2]+inc[2]=3+0=3
    data=[1,2], inc=[0,100], top=2
  5. top() → returns data[1]+inc[1]=2+100=102
  6. pop() →
    top-- → top=1
    inc[0] += inc[1] → inc[0]=0+100=100
    returns data[1]+inc[1]=2+100=102
    data=[1], inc=[100], top=1
  7. top() → returns data[0]+inc[0]=1+100=101

Through this example, you can see that the increment operation is correctly applied lazily when elements are accessed.