beam: trapping_rain_water two-pointer + prefix, tested

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

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