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.
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)