blob: e531a52ea7d00150d1b4e7eb7d959736c77bb015 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
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())
|