Design a Stack With Increment Operation

medium stack design lazy propagation

Problem

Design a CustomStack bounded by maxSize. Support push(x) (add to top only if not full), pop() (remove and return the top, or −1 if empty), and increment(k, val) which adds val to the bottom k elements (or all of them if fewer than k exist). The trick is to make increment run in O(1) instead of O(k).

Opspush 1, push 2, push 3, increment(2,100), pop, pop, pop

maxSize3

Stack [1,2,3]; increment bottom 2 by 100 → [101,102,3]; pops return 3, 102, 101.

maxSize

operations (e.g. push 5, pop, inc 2 100)

class CustomStack:
    def __init__(self, maxSize):
        self.n = maxSize
        self.stk = []        # element values
        self.inc = []        # lazy increment carried by each slot

    def push(self, x):
        if len(self.stk) < self.n:
            self.stk.append(x)
            self.inc.append(0)

    def pop(self):
        if not self.stk:
            return -1
        i = len(self.stk) - 1
        if i > 0:
            self.inc[i - 1] += self.inc[i]   # push pending down
        res = self.stk.pop() + self.inc.pop()
        return res

    def increment(self, k, val):
        i = min(k, len(self.stk)) - 1
        if i >= 0:
            self.inc[i] += val               # O(1) lazy bump

class CustomStack {
  constructor(maxSize) {
    this.n = maxSize;
    this.stk = [];          // element values
    this.inc = [];          // lazy increment per slot
  }
  push(x) {
    if (this.stk.length < this.n) {
      this.stk.push(x);
      this.inc.push(0);
    }
  }
  pop() {
    if (this.stk.length === 0) return -1;
    const i = this.stk.length - 1;
    if (i > 0) this.inc[i - 1] += this.inc[i];  // push pending down
    return this.stk.pop() + this.inc.pop();
  }
  increment(k, val) {
    const i = Math.min(k, this.stk.length) - 1;
    if (i >= 0) this.inc[i] += val;             // O(1) lazy bump
  }
}

class CustomStack {
    int n, size = 0;
    int[] stk, inc;            // values and lazy increments
    CustomStack(int maxSize) {
        n = maxSize;
        stk = new int[n];
        inc = new int[n];
    }
    void push(int x) {
        if (size < n) { stk[size] = x; inc[size] = 0; size++; }
    }
    int pop() {
        if (size == 0) return -1;
        int i = size - 1;
        if (i > 0) inc[i - 1] += inc[i];        // push pending down
        int res = stk[i] + inc[i];
        size--;
        return res;
    }
    void increment(int k, int val) {
        int i = Math.min(k, size) - 1;
        if (i >= 0) inc[i] += val;              // O(1) lazy bump
    }
}

class CustomStack {
    int n, size = 0;
    vector<int> stk, inc;      // values and lazy increments
public:
    CustomStack(int maxSize) : n(maxSize), stk(maxSize), inc(maxSize) {}
    void push(int x) {
        if (size < n) { stk[size] = x; inc[size] = 0; size++; }
    }
    int pop() {
        if (size == 0) return -1;
        int i = size - 1;
        if (i > 0) inc[i - 1] += inc[i];        // push pending down
        int res = stk[i] + inc[i];
        size--;
        return res;
    }
    void increment(int k, int val) {
        int i = min(k, size) - 1;
        if (i >= 0) inc[i] += val;              // O(1) lazy bump
    }
};

Explanation

The naive increment(k, val) walks the bottom k slots and adds val to each, costing O(k) per call. The clever version makes it O(1) by deferring the work with a parallel inc array — a form of lazy propagation.

Keep two arrays the same length as the stack: stk holds the values that were pushed, and inc holds a pending bonus for each slot. The true value at index i is stk[i] + inc[i], but we only materialise that bonus when an element is removed.

increment(k, val) needs to add val to the bottom k elements. Instead of touching all of them, we record it once at the top of that range: index i = min(k, size) − 1, doing inc[i] += val. The meaning is "everything at index i and below owes val".

pop() reads index i at the top. Before removing it we push its pending increment down one slot with inc[i-1] += inc[i], so the bonus reaches the elements below that still need it. Then the popped value is stk[i] + inc[i].

push(x) simply appends x with a fresh inc of 0, but only while the stack is below maxSize; otherwise the push is ignored. Each operation is now O(1).

Time: O(1) per op Space: O(maxSize)