Design a Stack With Increment Operation

Intuition

For each operation, my solution would only use O(1) time. Since each time we would only increase the bottom k values, we could just mark the k and when we need to pop, do the increase.

Solution


class CustomStack:
    def __init__(self, maxSize: int):
        self.stack = []
        self.size = maxSize
        self.increase = []

    def push(self, x: int) -> None:
        if len(self.stack) == self.size:
            return
        self.stack.append(x)
        self.increase.append(0)

    def pop(self) -> int:
        if not self.stack:
            return -1

        inc = self.increase.pop()
        if self.increase:
            self.increase[-1] += inc
        return inc + self.stack.pop()

    def increment(self, k: int, val: int) -> None:
        if not self.stack or k == 0:
            return

        if k > len(self.stack):
            self.increase[-1] += val
        else:
            self.increase[k - 1] += val


# Your CustomStack object will be instantiated and called as such:
# obj = CustomStack(maxSize)
# obj.push(x)
# param_2 = obj.pop()
# obj.increment(k,val)