60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
"""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())
|