aboutsummaryrefslogtreecommitdiff
path: root/design_stack_increment_op.py
diff options
context:
space:
mode:
Diffstat (limited to 'design_stack_increment_op.py')
-rw-r--r--design_stack_increment_op.py61
1 files changed, 61 insertions, 0 deletions
diff --git a/design_stack_increment_op.py b/design_stack_increment_op.py
new file mode 100644
index 0000000..e531a52
--- /dev/null
+++ b/design_stack_increment_op.py
@@ -0,0 +1,61 @@
+# 1381. Design a Stack With Increment Operation
+
+"""
+design a stack that supports increment operations on its elements. implement
+the function increment that increments the bottom 'k' elements of the stack
+by 'val' if there are less than 'k' elements in the stack, increment all the
+elements in the stack
+"""
+
+
+class CustomStack(object):
+ def __init__(self, maxSize):
+ """
+ :type maxSize: int
+ """
+ self.n = maxSize
+ self.stack = []
+ self.inc = []
+
+ def push(self, x):
+ """
+ :type x: int
+ :rtype: None
+ """
+ if len(self.inc) < self.n:
+ self.stack.append(x)
+ self.inc.append(0)
+
+ def pop(self):
+ """
+ :rtype: int
+ """
+ if not self.inc:
+ return -1
+ if len(self.inc) > 1:
+ self.inc[-2] += self.inc[-1]
+ return self.stack.pop() + self.inc.pop()
+
+ def increment(self, k, val):
+ """
+ :type k: int
+ :type val: int
+ :rtype: None
+ """
+ if self.inc:
+ self.inc[min(k, len(self.inc)) - 1] += val
+
+
+if __name__ == "__main__":
+ obj = CustomStack(3)
+ obj.push(1)
+ obj.push(2)
+ print(obj.pop())
+ obj.push(2)
+ obj.push(3)
+ obj.push(4)
+ obj.increment(5, 100)
+ print(obj.pop())
+ print(obj.pop())
+ print(obj.pop())
+ print(obj.pop())