Leetcode 1497 Solution

This article provides solution to leetcode question 1497 (design-a-stack-with-increment-operation)

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

Solution

class CustomStack:
def __init__(self, maxSize: int): self.maxSize = maxSize self.vals = [] self.inc_vals = collections.defaultdict(int)
def push(self, x: int) -> None: if len(self.vals) == self.maxSize: return
self.vals.append(x)
def pop(self) -> int: if not self.vals: return -1
index = len(self.vals) - 1
if index in self.inc_vals: inc = self.inc_vals[index] del self.inc_vals[index]
if index > 0: self.inc_vals[index - 1] += inc else: inc = 0
return self.vals.pop() + inc
def increment(self, k: int, val: int) -> None: k = min(k, len(self.vals)) self.inc_vals[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)