Coverage for blog/dsa/leetcode/stack_incr/__init__.py: 99%

1import pathlib 

2from typing import Literal 

3 

4import pytest 

5import yaml 

6 

7 

8# start snippet not_lazy 

9class CustomStackInitial: 

10 

11 top_max: int 

12 top: int 

13 data: list[int] 

14 

15 def __init__(self, maxSize: int): 

16 

17 self.top_max = maxSize - 1 

18 self.top = -1 

19 self.data = [-1 for _ in range(maxSize)] 

20 

21 def push(self, x: int) -> None: 

22 if self.top == self.top_max: 

23 return 

24 

25 self.top += 1 

26 self.data[self.top] = x 

27 

28 def pop(self) -> int: 

29 if self.top < 0: 

30 return -1 

31 

32 x = self.data[self.top] 

33 self.data[self.top] = -1 

34 self.top -= 1 

35 return x 

36 

37 def increment(self, k: int, val: int) -> None: 

38 

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 

43 

44 return 

45 # end snippet not_lazy 

46 

47 

48# start snippet solution 

49class CustomStack: 

50 

51 top_max: int 

52 top: int 

53 data: list[int] 

54 

55 def __init__(self, maxSize: int): 

56 

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

61 

62 def push(self, x: int) -> None: 

63 if self.top == self.top_max: 

64 return 

65 

66 self.top += 1 

67 self.data[self.top] = x 

68 

69 def pop(self) -> int: 

70 if self.top < 0: 

71 return -1 

72 

73 x = self.data[self.top] 

74 y = self.increments[self.top] 

75 

76 self.data[self.top] = -1 

77 self.increments[self.top] = 0 

78 self.top -= 1 

79 

80 if self.top > -1: 

81 self.increments[self.top] += y 

82 

83 return x + y 

84 

85 def increment(self, k: int, val: int) -> None: 

86 

87 if self.top < 0: 

88 return 

89 

90 index = min(k - 1, self.top) 

91 self.increments[index] += val 

92 

93 return 

94 # end snippet solution 

95 

96 

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

100 

101 

102def test_customstack_basic(): 

103 stack = CustomStackInitial(3) 

104 

105 assert stack.top == -1 

106 assert all(-1 == item for item in stack.data) 

107 assert stack.top_max == 2 

108 

109 stack.push(1) 

110 stack.push(2) 

111 

112 assert stack.top == 1 

113 assert stack.data == [1, 2, -1] 

114 

115 assert 2 == stack.pop() 

116 assert stack.top == 0 

117 assert stack.data == [1, -1, -1] 

118 

119 stack.push(2) 

120 stack.push(3) 

121 

122 assert stack.data == [1, 2, 3] 

123 assert stack.top == 2 

124 

125 stack.push(5) 

126 assert stack.data == [1, 2, 3] 

127 assert stack.top == 2 

128 

129 stack.increment(5, 100) 

130 assert stack.data == [101, 102, 103] 

131 

132 stack.increment(3, 100) 

133 assert stack.data == [201, 202, 203] 

134 

135 stack.increment(1, 100) 

136 assert stack.data == [301, 202, 203] 

137 

138 stack.increment(0, 666) 

139 assert stack.data == [301, 202, 203] 

140 

141 

142def test_customstack_basic_2(): 

143 stack = CustomStack(3) 

144 

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) 

148 

149 stack.push(1) 

150 stack.push(2) 

151 assert stack.data == [1, 2, -1] 

152 assert stack.increments == [0, 0, 0] 

153 

154 stack.push(3) 

155 assert stack.data == [1, 2, 3] 

156 assert stack.top == 2 

157 

158 stack.push(5) 

159 assert stack.data == [1, 2, 3] 

160 assert stack.top == 2 

161 

162 stack.increment(3, 5) 

163 assert stack.data == [1, 2, 3] 

164 assert stack.increments == [0, 0, 5] 

165 

166 assert stack.pop() == 8 

167 assert stack.data == [1, 2, -1] 

168 assert stack.increments == [0, 5, 0] 

169 

170 assert stack.pop() == 7 

171 assert stack.data == [1, -1, -1] 

172 assert stack.increments == [5, 0, 0] 

173 

174 stack.push(3) 

175 assert stack.data == [1, 3, -1] 

176 assert stack.increments == [5, 0, 0] 

177 

178 assert stack.pop() == 3 

179 assert stack.pop() == 6 

180 assert stack.pop() == -1 

181 

182 stack.push(55) 

183 assert stack.data == [55, -1, -1] 

184 assert stack.increments == [0, 0, 0] 

185 

186 stack.push(66) 

187 stack.push(77) 

188 assert stack.data == [55, 66, 77] 

189 

190 stack.increment(5, 100) 

191 assert stack.increments == [0, 0, 100] 

192 

193 stack.increment(2, 100) 

194 assert stack.increments == [0, 100, 100] 

195 

196 assert stack.pop() == 177 

197 assert stack.increments == [0, 200, 0] 

198 

199 assert stack.pop() == 266 

200 assert stack.increments == [200, 0, 0] 

201 

202 assert stack.pop() == 255 

203 assert stack.increments == [0, 0, 0] 

204 

205 

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 = [] 

215 

216 for method_name, method_args in zip(methods, args): 

217 

218 if (method := getattr(stack, method_name, None)) is None: 

219 raise ValueError(f"No such method `{method_name}` of `Stack`.") 

220 

221 result_computed.append(method(*method_args)) 

222 

223 print("-------------------------------------------------") 

224 print(method_name, method_args) 

225 print(list(zip(stack.increments, stack.data))) 

226 print(stack.top) 

227 

228 print(result) 

229 print(result_computed) 

230 assert result == result_computed