LeetCode 42. Trapping Rain Water

October 05, 2020

hardLeetCode

We can do this in $O(n^2)$ time:

For each bar:

Backtrack in the array and find the tallest prior bar that has height <= the current bar (or the first one that's at least the same height).
Add $\lt curr\_index \gt - \lt bar\_from\_step\_1\_index \gt$ units to the cumulative trapped water.

I'm guessing we can do this in linear time

Can we use a stack to keep the relevant prior bars in memory?

For each bar:

Remove bars from the stack if bar is shorter than current (add water: distance * [height of the shorter bar - max height of stack-bars already seen])

Stop when you reach a bar that's at least the same height (add water: same as above)

Put the current bar onto the stack.

from typing import List
import unittest

class Test(unittest.TestCase):

    def setUp(self):
        self.s = Solution()

    def test_examples(self):
        self.assertEqual(self.s.trap([0,1,0,2,1,0,1,3,2,1,2,1]), 6)

    def test_empty(self):
        self.assertEqual(self.s.trap([]), 0)
        self.assertEqual(self.s.trap([0]), 0)
        self.assertEqual(self.s.trap([0,0]), 0)
        self.assertEqual(self.s.trap([0,0,0]), 0)
        self.assertEqual(self.s.trap([0,0,0,1]), 0)
    
    def test_tall_first(self):
        self.assertEqual(self.s.trap([3,0,2,1]), 2)
        self.assertEqual(self.s.trap([0,3,0,2,1]), 2)
        self.assertEqual(self.s.trap([0,3,0,2,1,0]), 2)

    def test_zero_middle(self):
        self.assertEqual(self.s.trap([3,0,0,0,3]), 9)

class Solution:

    def trap(self, height: List[int]) -> int:
        water = 0
        priors = []
        for i in range(len(height)):
            max_seen = 0
            b = height[i]
            while b > 0 and priors:
                p = priors[-1]
                water += (i - p[1] - 1) * (min(p[0], b) - max_seen)
                max_seen = p[0]
                if p[0] > b:
                    break
                priors.pop()
            priors.append((b,i))
        return water

if __name__ == "__main__":
    unittest.main()

LeetCode 42. Trapping Rain Water

We can do this in O(n2)O(n^2)O(n2) time:

I'm guessing we can do this in linear time

We can do this in $O(n^2)$ time: