diff --git a/LeetCode/hard/trap_42.py b/LeetCode/hard/trap_42.py new file mode 100644 index 0000000..23e8785 --- /dev/null +++ b/LeetCode/hard/trap_42.py @@ -0,0 +1,23 @@ +from typing import List + + +class Solution: + def trap(self, height: List[int]) -> int: + if not height: + return 0 + + left, right = 0, len(height) - 1 + left_max, right_max = height[left], height[right] + res = 0 + + while left < right: + if left_max < right_max: + left += 1 + left_max = max(left_max, height[left]) + res += left_max - height[left] + else: + right -= 1 + right_max = max(right_max, height[right]) + res += right_max - height[right] + + return res diff --git a/tests/test_leecode_hard.py b/tests/test_leecode_hard.py new file mode 100644 index 0000000..65a1940 --- /dev/null +++ b/tests/test_leecode_hard.py @@ -0,0 +1,23 @@ +import pytest + + +@pytest.mark.parametrize( + "height, expected", + [ + ([0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1], 6), + ([4, 2, 0, 3, 2, 5], 9), + ([], 0), + ([1], 0), + ([1, 2], 0), + ([1, 2, 3, 4, 5], 0), # increasing + ([5, 4, 3, 2, 1], 0), # decreasing + ([3, 3, 3, 3], 0), # flat + ([3, 0, 3], 3), + ([5, 0, 5], 5), + ([2, 0, 2], 2), + ], +) +def test_trap(height, expected): + from LeetCode.hard.trap_42 import Solution + + assert Solution().trap(height) == expected