Design a Stack With Increment Operation: Lazy Propagation Pattern
You are asked to design a custom stack that supports three operations: push(x) adds an element to the top, pop() removes and returns the top element (or -1 if empty), and increment(k, val) adds val to the bottom k elements of the stack. The stack has a fixed maximum size.
This is LeetCode 1381: Design a Stack With Increment Operation, and it is one of the best problems for learning the lazy propagation pattern. The naive approach updates k elements on every increment call, but the clever approach defers that work until it is actually needed.
Why this problem matters
At first glance, the increment operation seems to require O(k) work every time it is called. You walk through the bottom k elements and add val to each one. That works, but if increment is called frequently with large k, the cost adds up fast.
The key idea is that you do not need to apply the increment immediately. You can record it and apply it later, only when an element is actually popped. This is the same "lazy" principle that shows up in segment trees, lazy deletion in heaps, and deferred computation throughout computer science. Learning it here, in a simple stack context, makes it easier to recognize and apply in harder problems.
The key insight
Maintain an auxiliary array inc of the same size as the stack. When increment(k, val) is called, instead of updating k elements, just add val to inc[min(k, size) - 1]. This single write replaces an O(k) loop.
When you pop an element, add inc[top] to the element's value before returning it. Then propagate the increment downward: if there is an element below, add inc[top] to inc[top - 1]. This ensures the increment eventually reaches every element it was meant for, but only when each element is actually accessed.
The algorithm:
- push(x): If the stack is not full, push
xand setinc[top] = 0. - pop(): If the stack is empty, return
-1. Otherwise, getresult = stack[top] + inc[top]. Iftop > 0, propagateinc[top]down toinc[top - 1]. Resetinc[top] = 0, decrement top, and returnresult. - increment(k, val): Set
idx = min(k, size) - 1. Addvaltoinc[idx].
The solution
class CustomStack:
def __init__(self, maxSize: int):
self.stack = []
self.inc = []
self.max_size = maxSize
def push(self, x: int) -> None:
if len(self.stack) < self.max_size:
self.stack.append(x)
self.inc.append(0)
def pop(self) -> int:
if not self.stack:
return -1
idx = len(self.stack) - 1
result = self.stack[idx] + self.inc[idx]
if idx > 0:
self.inc[idx - 1] += self.inc[idx]
self.stack.pop()
self.inc.pop()
return result
def increment(self, k: int, val: int) -> None:
if self.stack:
idx = min(k, len(self.stack)) - 1
self.inc[idx] += val
Let's walk through what each piece does.
The constructor initializes two lists: stack for the actual values and inc for the lazy increments. Both grow and shrink together, staying the same length at all times.
The push method is simple. If there is room, append the value and a corresponding 0 in the increment array. The 0 means "no pending increment for this position yet."
The pop method is where the lazy magic happens. The return value is stack[idx] + inc[idx], combining the stored value with any accumulated increment. Before removing the element, if there is an element below it, the current increment gets added to the element below. This propagation step is critical. It ensures that increments meant for multiple elements trickle down as elements are removed, one at a time.
The increment method does the minimum possible work. It finds the topmost index that needs the increment (the smaller of k and the current stack size, minus one) and adds val there. Everything below that index will eventually receive the increment through propagation during pops.
The propagation in pop is the core trick. When you pop index i, any increment stored at inc[i] gets added to inc[i - 1]. This means increments "sink" downward one level at a time, reaching each element exactly when it becomes the top. You never touch more than one element per pop.
Visual walkthrough
Let's trace through a sequence of operations to see how the lazy increment array works in practice. Watch how the increment is stored once and then propagated during pops.
Step 1: push(1), push(2), push(3)
Stack has 3 elements. All inc values are 0.
Step 2: increment(2, 100) — add 100 to bottom 2 elements
Instead of updating stack[0] and stack[1], we just set inc[1] = 100. Lazy!
Step 3: pop() — remove top element (3)
Pop returns stack[2] + inc[2] = 3 + 0 = 3. inc[2] was 0, so nothing to propagate.
Step 4: pop() — remove top element, apply lazy increment
Pop returns stack[1] + inc[1] = 2 + 100 = 102. Before removing, propagate: inc[0] += inc[1], so inc[0] = 0 + 100 = 100.
Step 5: pop() — remove last element, apply propagated increment
Pop returns stack[0] + inc[0] = 1 + 100 = 101. The increment originally set at index 1 has propagated all the way down.
Step 6: pop() — stack is empty
Stack is empty. Return -1.
The key moment is Steps 4 and 5. When we pop the element at index 1, its increment (100) propagates down to index 0. Then when we pop index 0, the propagated increment is applied, giving us 1 + 100 = 101. The increment was set once at index 1 but correctly affected both elements below it.
Complexity analysis
| Approach | push | pop | increment |
|---|---|---|---|
| Naive (update k elements) | O(1) | O(1) | O(k) |
| Lazy increment array | O(1) | O(1) | O(1) |
Push is O(1) in both approaches. You append to the end of the stack.
Pop is O(1) in both approaches. With the lazy approach, you do a constant amount of extra work: read inc[top], propagate to inc[top - 1], and reset.
Increment is where the difference matters. The naive approach loops through min(k, size) elements. The lazy approach writes to a single index, making it O(1). Over a sequence of n operations, the lazy approach saves significant work.
Space is O(n) for both approaches, where n is the maximum stack size. The lazy approach uses an extra array of the same size as the stack.
The building blocks
1. Lazy propagation on a stack
The pattern of deferring updates and applying them only when needed:
def pop(self) -> int:
idx = len(self.stack) - 1
result = self.stack[idx] + self.inc[idx]
if idx > 0:
self.inc[idx - 1] += self.inc[idx]
self.stack.pop()
self.inc.pop()
return result
This is a simplified version of the lazy propagation used in segment trees and other advanced data structures. The principle is the same: store the update, propagate it only when you need to access the affected data. You will see this pattern in range update queries, lazy segment trees, and any scenario where batch updates can be deferred.
2. Auxiliary array mirroring a data structure
The pattern of maintaining a parallel array that tracks extra metadata for each element:
self.stack = []
self.inc = []
Both arrays grow and shrink together. Each index in inc corresponds to the same index in stack. This "shadow array" pattern appears in many design problems. The Min Stack problem uses a parallel stack to track minimums. Monotonic stack problems sometimes use auxiliary arrays to track indices or counts. The idea of pairing each element with supplementary information is a fundamental design tool.
3. Targeted index update for range operations
The pattern of touching a single index to represent an operation on a range:
def increment(self, k: int, val: int) -> None:
idx = min(k, len(self.stack)) - 1
self.inc[idx] += val
Instead of updating indices 0 through k-1, you store the entire update at one boundary index. This is similar to difference arrays, where you mark the start and end of a range update and reconstruct actual values later. The single-write approach turns O(k) operations into O(1).
Edge cases
Before submitting, think through these scenarios:
- Stack is full:
pushshould be a no-op when the stack has reachedmaxSize. - Pop on empty stack: return
-1. Make sure you check before accessingstack[idx]. - Increment with k larger than stack size: use
min(k, len(stack))so you do not index out of bounds. - Increment on empty stack: the
if self.stackguard prevents an index error when the stack is empty. - Multiple increments before any pop: they accumulate at the same or different indices. The propagation during pops correctly combines them.
- Single element stack: increment and pop both work. There is no element below to propagate to, and the
if idx > 0check handles this.
From understanding to recall
You have seen how lazy propagation turns an O(k) increment into O(1) by deferring the work to pop time. The logic is clean and the code is short. But reproducing it from memory under pressure is a different challenge.
The details that trip people up are subtle. Do you propagate before or after popping? Do you use min(k, len(stack)) or min(k, len(stack)) - 1? Do you reset inc[top] after propagating? These are not conceptual gaps. They are recall gaps, and they cost people offers.
Spaced repetition is how you close them. You practice writing the lazy increment setup, the propagation logic in pop, and the single-index write in increment from memory at increasing intervals. After a few rounds, the pattern is automatic. You see "design a data structure with deferred updates" in a problem description, and the lazy propagation template flows out without hesitation.
Related posts
- Min Stack - Another stack design problem that uses an auxiliary data structure to track extra information per element
- Implement Queue Using Stacks - Stack design problem that teaches you to think about amortized costs across operations
- Daily Temperatures - Uses a monotonic stack with index tracking, another example of pairing stack elements with metadata
CodeBricks breaks Design a Stack With Increment Operation into its lazy propagation and auxiliary array building blocks, then drills them independently with spaced repetition. You type each piece from scratch until the pattern is automatic. When a design problem shows up in your interview, you do not think about it. You just write it.