Stack With a Range-Increment Operation

Medium25–30 min

Topics

StackDesign

Companies

AmazonAdobeUber

Design a stack that supports the usual push and pop, plus increment(k, val), which adds val to each of the bottom min(k, size) elements currently on the stack. Applying an increment to potentially many elements at once seems to demand touching each of them — but the actual values only ever matter once an element is popped and revealed. That gap between "when the increment happens" and "when the value is needed" is exploitable: record each increment's effect in exactly one place (a parallel array, at the position marking how far down it reaches), and let that pending amount quietly cascade downward one step at a time, only when a pop makes it necessary. Every operation — push, pop, and increment alike — becomes a fixed, small amount of work, no matter how many elements an increment conceptually covers.

Test Case 1:

Input:ops = ["push","push","push","increment","pop","pop","pop"], args = [4,3,2,2,0,0,0], args2 = [0,0,0,10,0,0,0]

Output:[2, 13, 14]

Explanation:After pushing 4,3,2 (bottom to top), increment(2,10) adds 10 to the bottom 2 elements (4 and 3), making the stack 14,13,2 (bottom to top). Popping unwinds it top to bottom: 2, 13, 14.

Test Case 2:

Input:ops = ["push","push","increment","push","pop","pop","pop"], args = [1,2,100,3,0,0,0], args2 = [0,0,5,0,0,0,0]

Output:[3, 7, 6]

Explanation:increment(100,5) with k larger than the stack's size (2) just adds 5 to every element that exists. After pushing 3 on top, popping gives 3, then 7 (=2+5), then 6 (=1+5).

Test Case 3:

Input:ops = ["push","push","push","increment","increment","pop","pop","pop"], args = [5,5,5,1,2,0,0,0], args2 = [0,0,0,3,4,0,0,0]

Output:[5, 9, 12]

Explanation:increment(1,3) adds 3 only to the very bottom element. increment(2,4) then adds 4 to the bottom 2 elements. The increments stack up: bottom becomes 5+3+4=12, second becomes 5+4=9, top stays 5.

Constraints

◆1 ≤ ops.length ≤ 20
◆ops[i] is one of "push", "pop", "increment"; args[i] is the pushed value (for push) or k (for increment); args2[i] is val (for increment), 0 otherwise
◆increment(k, val) adds val to each of the bottom min(k, current size) elements
◆"pop" on an empty stack returns -1

🚀

Try the Dry Run

Don't just read the solution — watch it execute, one step at a time.

Approach & Solutions

Brute Force — Directly Add to Every Affected Element

Brute

Use a plain array-backed stack. push and pop are the usual O(1) operations. For increment(k, val), walk the bottom min(k, size) elements and add val to each one directly. This is completely correct and simple, but every increment call costs time proportional to how many elements it touches — a sequence with many large-k increments pays that cost every single time.

TimeO(1) push/pop, O(k) increment

SpaceO(n)

Java

Try it out

class Solution {
    public int[] stackOpsIncrement(String[] ops, int[] args, int[] args2) {
        List<Integer> stack = new ArrayList<>();
        List<Integer> results = new ArrayList<>();
        for (int i = 0; i < ops.length; i++) {
            String op = ops[i];
            if (op.equals("push")) {
                stack.add(args[i]);
            } else if (op.equals("pop")) {
                results.add(stack.isEmpty() ? -1 : stack.remove(stack.size() - 1));
            } else {
                int k = args[i], val = args2[i];
                int limit = Math.min(k, stack.size());
                for (int j = 0; j < limit; j++) {
                    stack.set(j, stack.get(j) + val);
                }
            }
        }
        int[] output = new int[results.size()];
        for (int i = 0; i < results.size(); i++) output[i] = results.get(i);
        return output;
    }
}

Optimal — Lazy Increments Resolved Only at Pop Time

Optimal

Never touch min(k, size) elements up front. Instead, keep a second array of pending increments, one slot per stack position. increment(k, val) records its effect in exactly one place — the slot at index k-1 (the boundary of what it affects) — in O(1), without touching anything else. The trick is deferring the actual work: a pending increment at some position secretly also applies to everything below it, so when popping the top element, first push that position's pending amount down onto the position right below it (so it isn't lost), then add the (now-correct) pending amount to the value being returned. Every push, pop, and increment call does a fixed, small amount of work — O(1) across the board.

TimeO(1) for every operation