Problem statement

https://leetcode.com/problems/design-a-stack-with-increment-operation/

Solution

The idea is to keep differences of adjacent elements in our stack as well as the value of last element. In this way we can make all operations in O(1).

Complexity

Time complexity of all operations is O(1), space is O(n).

Code

class CustomStack:
    def __init__(self, maxSize):
        self.stack = [0]
        self.size = maxSize
        self.last = 0
        
    def push(self, x):
        if len(self.stack) <= self.size:
            self.stack += [x - self.last]
            self.last = x
        
    def pop(self):
        if len(self.stack) != 1:
            ans = self.last
            self.last -= self.stack.pop()
            return ans
        return -1
        
    def increment(self, k, val):
        if k < len(self.stack) - 1:
            self.stack[1] += val
            self.stack[k + 1] -= val
        elif len(self.stack) > 1:
            self.stack[1] += val
            self.last += val