leetcode

Leetcode Solutions (0x6A73)
git clone https://git.pngu.org/leetcode
Log | Files | Refs | Submodules | README

design_stack_increment_op.py (1377B)

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