I'm not sure what you mean by "used too much memory" — if I try submitting a copy of your code now, then LeetCode says:
Memory Usage: 12.3 MB, less than 57.53% of Python online submissions for Trapping Rain Water.
It is impossible to know whether the 12.3 MB is a lot or a little without knowing what this includes and how big the test cases were, neither of which LeetCode tells us.
I found it difficult to follow the algorithm because of the proliferation of similarly named variables. There are four variables that contain heights:
max_left_bound. How do these variables relate to each other? Is
left_bound actually a right bound when iterating over heights in reverse? There are three variables that contain amounts of water:
block_capacity. Again, how do these relate to each other?
Finding good names for variables and writing comments to clarify any remaining difficulties, is crucial in writing code that will be maintainable when you come back to it after having forgotten how it works.
trap function modifies the list
heights (by reversing it). It's best if functions don't modify their arguments (unless modifying the arguments is the point). Otherwise callers will be surprised by the side-effects. In this case you could easily use
reversed(heights) instead of
The code considers each vertical slice through the landscape in turn, adding
left_bound - elevation (that is, the current water height minus the current land height) to the total trapped water. But in practice
left_bound (the current water height) stays the same for several iterations of the loop, so we could wait until it changes and then add the block of water all at once. See §2 below for how to do this.
2. Alternative solution
Consider a typical landscape filled with water:
If we ignore the difference between water and land, and just look at the overall surface level, it should be clear that the filled landscape is made up of rectangles arranged like this:
There is a central "plateau" whose height is the maximum height of the landscape, flanked by staircases that start low at the outside and rise up to the plateau.
If we find the height of the plateau, which is just
max(landscape), then we can work inwards from the left and right ends, finding the rectangles in the rising staircases until we reach the plateau height. By summing the rectangles we get the overall area of the filled landscape, and by subtracting the land part, which is just
sum(landscape), we get the water part.
"""Return area of water trapped by landscape (a list of heights)."""
if not landscape:
plateau_height = max(landscape)
plateau_width = len(landscape)
total_area = 0
for heights in (landscape, reversed(landscape)):
height_max = i_max = 0 # maximum height so far and its index
for i, height in enumerate(heights):
if height > height_max:
total_area += (i - i_max) * height_max
if height == plateau_height:
plateau_width -= i
height_max = height
i_max = i
return total_area + plateau_height * plateau_width - sum(landscape)
It is a bit clearer, I hope, how this works, and it is several times faster than the code in the post. The reasons why this is faster are (i) the accumulation
total_area += ... is done once per rectangle, not once per vertical slice; and (ii) the subtraction of the land part is done once by calling the built-in function
sum, which runs in fast compiled code.
The exact speedup depends on the shape of the landscape, but here's a test case where the revised code was about 8 times faster:
>>> from random import choices
>>> from timeit import timeit
>>> HEIGHTS = choices(range(100), weights=(100,) + (1,) * 99, k=1000000)
>>> timeit(lambda:trap(HEIGHTS), number=1) # code in the post
>>> timeit(lambda:trap2(HEIGHTS), number=1) # revised code
But on LeetCode it runs in much the same time as the code in the post — clearly they are not try big enough test cases to detect the difference.