-
Notifications
You must be signed in to change notification settings - Fork 4
Expand file tree
/
Copy pathtrappingRainWater.py
More file actions
29 lines (23 loc) · 826 Bytes
/
trappingRainWater.py
File metadata and controls
29 lines (23 loc) · 826 Bytes
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
import unittest
from typing import List
class Solution:
def trap(self, height: List[int]) -> int:
stack = []
volume = 0
for i in range(len(height)):
while stack and height[i] > height[stack[-1]]:
top = stack.pop()
if not len(stack):
break
distance = i - stack[-1] - 1
waters = min(height[i], height[stack[-1]]) - height[top]
volume += distance * waters
stack.append(i)
return volume
class Test(unittest.TestCase):
def test_trap(self):
solution = Solution()
self.assertEqual(solution.trap([0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]), 6)
self.assertEqual(solution.trap([4, 2, 0, 3, 2, 5]), 9)
if __name__ == '__main__':
unittest.main()