From 8e08516960d980d2f3a550770fb9db4a41f8125e Mon Sep 17 00:00:00 2001 From: beam Date: Mon, 2 Nov 2026 09:47:49 +0000 Subject: [PATCH] beam: trapping_rain_water two-pointer + prefix, tested --- wizard-work/beam/test_trapping_rain_water.py | 40 +++++++++++++ wizard-work/beam/trapping_rain_water.py | 59 ++++++++++++++++++++ 2 files changed, 99 insertions(+) create mode 100644 wizard-work/beam/test_trapping_rain_water.py create mode 100644 wizard-work/beam/trapping_rain_water.py diff --git a/wizard-work/beam/test_trapping_rain_water.py b/wizard-work/beam/test_trapping_rain_water.py new file mode 100644 index 0000000..46cafe2 --- /dev/null +++ b/wizard-work/beam/test_trapping_rain_water.py @@ -0,0 +1,40 @@ +import unittest +from trapping_rain_water import trap, trap_prefix + +class TestTrap(unittest.TestCase): + def test_empty(self): + self.assertEqual(trap([]), 0) + self.assertEqual(trap_prefix([]), 0) + + def test_single(self): + self.assertEqual(trap([5]), 0) + + def test_flat(self): + self.assertEqual(trap([2, 2, 2, 2]), 0) + + def test_classic(self): + self.assertEqual(trap([0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]), 6) + + def test_ascending(self): + self.assertEqual(trap([1, 2, 3, 4]), 0) + + def test_descending(self): + self.assertEqual(trap([4, 3, 2, 1]), 0) + + def test_two_peaks(self): + self.assertEqual(trap([3, 0, 2, 0, 4]), 7) + + def test_agreement_random(self): + import random + rng = random.Random(42) + for _ in range(2000): + n = rng.randint(0, 40) + hs = [rng.randint(0, 9) for _ in range(n)] + self.assertEqual(trap(hs), trap_prefix(hs)) + + def test_large(self): + hs = [10**6] + [0] * 100000 + [10**6] + self.assertEqual(trap(hs), 10**6 * 100000) + +if __name__ == "__main__": + unittest.main() diff --git a/wizard-work/beam/trapping_rain_water.py b/wizard-work/beam/trapping_rain_water.py new file mode 100644 index 0000000..943f508 --- /dev/null +++ b/wizard-work/beam/trapping_rain_water.py @@ -0,0 +1,59 @@ +"""Trapping rain water on an elevation map. + +Given heights[i] = elevation at bar i, compute total water trapped +after raining. Water trapped at i = max(0, min(max_left, max_right) - height[i]). + +Two implementations: O(n) two-pointer (trap) and O(n) prefix arrays (trap_prefix). +Both agree; trap is the production path. +""" +from typing import Sequence + +def trap(heights: Sequence[int]) -> int: + if not heights: + return 0 + left, right = 0, len(heights) - 1 + left_max, right_max = 0, 0 + water = 0 + while left < right: + if heights[left] < heights[right]: + if heights[left] >= left_max: + left_max = heights[left] + else: + water += left_max - heights[left] + left += 1 + else: + if heights[right] >= right_max: + right_max = heights[right] + else: + water += right_max - heights[right] + right -= 1 + return water + +def trap_prefix(heights: Sequence[int]) -> int: + n = len(heights) + if n == 0: + return 0 + left_max = [0] * n + right_max = [0] * n + cur = 0 + for i in range(n): + if heights[i] > cur: + cur = heights[i] + left_max[i] = cur + cur = 0 + for i in range(n - 1, -1, -1): + if heights[i] > cur: + cur = heights[i] + right_max[i] = cur + return sum(max(0, min(left_max[i], right_max[i]) - heights[i]) for i in range(n)) + +def main(): + import sys + if len(sys.argv) < 2: + print("usage: trapping_rain_water.py h1 h2 ..."); return 2 + hs = [int(x) for x in sys.argv[1:]] + print(trap(hs)) + return 0 + +if __name__ == "__main__": + raise SystemExit(main())