diff --git a/dynamic_programming/trapped_water.py b/dynamic_programming/trapped_water.py index 8bec9fac5fef..5d950efd7264 100644 --- a/dynamic_programming/trapped_water.py +++ b/dynamic_programming/trapped_water.py @@ -34,22 +34,23 @@ def trapped_rainwater(heights: tuple[int, ...]) -> int: return 0 if any(h < 0 for h in heights): raise ValueError("No height can be negative") - length = len(heights) - - left_max = [0] * length - left_max[0] = heights[0] - for i, height in enumerate(heights[1:], start=1): - left_max[i] = max(height, left_max[i - 1]) - - right_max = [0] * length - right_max[-1] = heights[-1] - for i in range(length - 2, -1, -1): - right_max[i] = max(heights[i], right_max[i + 1]) - - return sum( - min(left, right) - height - for left, right, height in zip(left_max, right_max, heights) - ) + left, right = 0, len(heights) - 1 + leftmax, rightmax = 0, 0 + water = 0 + while left < right: + if heights[left] < heights[right]: + if heights[left] >= leftmax: + leftmax = heights[left] + else: + water += leftmax - heights[left] + left += 1 + else: + if heights[right] >= rightmax: + rightmax = heights[right] + else: + water += rightmax - heights[right] + right -= 1 + return water if __name__ == "__main__":