aboutsummaryrefslogtreecommitdiff
path: root/design_stack_increment_op.py
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())