Coverage for blog/dsa/leetcode/stack_incr/__init__.py: 99%
150 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-20 16:23 +0000
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-20 16:23 +0000
1import pathlib
2from typing import Literal
4import pytest
5import yaml
8# start snippet not_lazy
9class CustomStackInitial:
11 top_max: int
12 top: int
13 data: list[int]
15 def __init__(self, maxSize: int):
17 self.top_max = maxSize - 1
18 self.top = -1
19 self.data = [-1 for _ in range(maxSize)]
21 def push(self, x: int) -> None:
22 if self.top == self.top_max:
23 return
25 self.top += 1
26 self.data[self.top] = x
28 def pop(self) -> int:
29 if self.top < 0:
30 return -1
32 x = self.data[self.top]
33 self.data[self.top] = -1
34 self.top -= 1
35 return x
37 def increment(self, k: int, val: int) -> None:
39 # NOTE: Use the minimum since if ``k`` exceeds ``top_max`` all elements
40 # should be incremented. Should do nothing when ``k=0``.
41 for j in range(min(k, self.top_max + 1)):
42 self.data[j] += val
44 return
45 # end snippet not_lazy
48# start snippet solution
49class CustomStack:
51 top_max: int
52 top: int
53 data: list[int]
55 def __init__(self, maxSize: int):
57 self.top_max = maxSize - 1
58 self.top = -1
59 self.data = [-1 for _ in range(maxSize)]
60 self.increments = [0 for _ in range(maxSize)]
62 def push(self, x: int) -> None:
63 if self.top == self.top_max:
64 return
66 self.top += 1
67 self.data[self.top] = x
69 def pop(self) -> int:
70 if self.top < 0:
71 return -1
73 x = self.data[self.top]
74 y = self.increments[self.top]
76 self.data[self.top] = -1
77 self.increments[self.top] = 0
78 self.top -= 1
80 if self.top > -1:
81 self.increments[self.top] += y
83 return x + y
85 def increment(self, k: int, val: int) -> None:
87 if self.top < 0:
88 return
90 index = min(k - 1, self.top)
91 self.increments[index] += val
93 return
94 # end snippet solution
97CASES = pathlib.Path(__file__).parent.resolve() / "cases.yaml"
98with open(CASES, "r") as file:
99 cases = tuple(item.values() for item in yaml.safe_load(file))
102def test_customstack_basic():
103 stack = CustomStackInitial(3)
105 assert stack.top == -1
106 assert all(-1 == item for item in stack.data)
107 assert stack.top_max == 2
109 stack.push(1)
110 stack.push(2)
112 assert stack.top == 1
113 assert stack.data == [1, 2, -1]
115 assert 2 == stack.pop()
116 assert stack.top == 0
117 assert stack.data == [1, -1, -1]
119 stack.push(2)
120 stack.push(3)
122 assert stack.data == [1, 2, 3]
123 assert stack.top == 2
125 stack.push(5)
126 assert stack.data == [1, 2, 3]
127 assert stack.top == 2
129 stack.increment(5, 100)
130 assert stack.data == [101, 102, 103]
132 stack.increment(3, 100)
133 assert stack.data == [201, 202, 203]
135 stack.increment(1, 100)
136 assert stack.data == [301, 202, 203]
138 stack.increment(0, 666)
139 assert stack.data == [301, 202, 203]
142def test_customstack_basic_2():
143 stack = CustomStack(3)
145 assert stack.top == -1 and stack.top_max == 2
146 assert all(-1 == item for item in stack.data)
147 assert all(0 == item for item in stack.increments)
149 stack.push(1)
150 stack.push(2)
151 assert stack.data == [1, 2, -1]
152 assert stack.increments == [0, 0, 0]
154 stack.push(3)
155 assert stack.data == [1, 2, 3]
156 assert stack.top == 2
158 stack.push(5)
159 assert stack.data == [1, 2, 3]
160 assert stack.top == 2
162 stack.increment(3, 5)
163 assert stack.data == [1, 2, 3]
164 assert stack.increments == [0, 0, 5]
166 assert stack.pop() == 8
167 assert stack.data == [1, 2, -1]
168 assert stack.increments == [0, 5, 0]
170 assert stack.pop() == 7
171 assert stack.data == [1, -1, -1]
172 assert stack.increments == [5, 0, 0]
174 stack.push(3)
175 assert stack.data == [1, 3, -1]
176 assert stack.increments == [5, 0, 0]
178 assert stack.pop() == 3
179 assert stack.pop() == 6
180 assert stack.pop() == -1
182 stack.push(55)
183 assert stack.data == [55, -1, -1]
184 assert stack.increments == [0, 0, 0]
186 stack.push(66)
187 stack.push(77)
188 assert stack.data == [55, 66, 77]
190 stack.increment(5, 100)
191 assert stack.increments == [0, 0, 100]
193 stack.increment(2, 100)
194 assert stack.increments == [0, 100, 100]
196 assert stack.pop() == 177
197 assert stack.increments == [0, 200, 0]
199 assert stack.pop() == 266
200 assert stack.increments == [200, 0, 0]
202 assert stack.pop() == 255
203 assert stack.increments == [0, 0, 0]
206@pytest.mark.parametrize("size, methods, args, result", cases)
207def test_solution(
208 size: int,
209 methods: list[Literal["push", "pop", "increment"]],
210 args: list[tuple[int] | tuple[int, int] | tuple],
211 result: list[None | int],
212):
213 stack = CustomStack(size)
214 result_computed = []
216 for method_name, method_args in zip(methods, args):
218 if (method := getattr(stack, method_name, None)) is None:
219 raise ValueError(f"No such method `{method_name}` of `Stack`.")
221 result_computed.append(method(*method_args))
223 print("-------------------------------------------------")
224 print(method_name, method_args)
225 print(list(zip(stack.increments, stack.data)))
226 print(stack.top)
228 print(result)
229 print(result_computed)
230 assert result == result_computed