Problem Statement in English

You’re given an array of integers nums, indicating the height of the walls. Your task is to return the total water that can be trapped in between all the walls put together.


Approach

Full disclosure: this is not my code, I got this online from NeetCode (brilliant channel by the way, definitely check him out).

There’s 2 approaches to this, one with $O(n)$ space complexity, and another with $O(1)$ space complexity.

  • Linear Space Solution: This uses a Prefix Sums based approach. We make a left pass and calculate the heightest wall up until the current index we’re processing, and a right pass which does the same thing, but from the right.

    Then for any given index we know the highest wall on its left and right. Using this we can calculate the amount of water that can be held at that index.

  • Constant Space Solution: This solution uses the 2 Pointer Technique. Check out this post as this question is a slight variation of that one.

    The objective is the same as in the linear space solution. We want to always have the highest walls on the left and right of any index that we’re processing. Just that here we don’t store these heights in arrays, hence saving on space.

Take a look at the code now.


Solution in Python

O(n) space


# O(n) space
class Solution:
    def trap(self, height: List[int]) -> int:
        # holds the final answer
        total = 0
        # to hold the current max height seen so far
        currBest = 0

        # holds result of the left pass
        leftHighest = [0 for _ in range(len(height))]
        # holds result of the right pass
        rightHighest = leftHighest.copy()

        # left pass
        for i in range(len(height)):
            leftHighest[i] = currBest
            currBest = max(currBest, height[i])

        # right pass
        currBest = 0
        for i in range(len(height) - 1, -1, -1):
            rightHighest[i] = currBest
            currBest = max(currBest, height[i])

        # compute result
        for i in range(len(height)):
            # calculate amount of water that can be trapped at that index
            trappable = min(leftHighest[i], rightHighest[i]) - height[i]
            # if it's non zero, add it to the solution
            if trappable > 0:
                total += trappable

        # return final answer
        return total

O(1) space


class Solution:
    def trap(self, height: List[int]) -> int:
        # holds the final answer
        total = 0

        # holds the value of the variable pointed to by the pointer that we moved last
        temp = 0

        # initialise left and right pointers
        l, r = 0, len(height) - 1

        # holds current best left and right wall height
        leftMax, rightMax = height[l], height[r]

        # run code while left pointer is before the right one
        while l < r:
            # if the left wall is smaller than the right wall
            if height[l] < height[r]:
                # we want it to be bigger to maximise water that can be held
                # so move it towards the right,
                # as we've already processed walls on the left of this one
                l += 1
                # update the value pointed by the last moved pointer
                temp = height[l]
            # if the right wall is smaller than the left wall
            else:
                # we want it to be bigger to maximise water that can be held
                # so move it towards the left,
                # as we've already processed walls on the right of this one
                r -= 1
                # update the value pointed by the last moved pointer
                temp = height[r]

            # depth of the container depends on the smaller of the 2 walls
            depth = min(leftMax, rightMax)

            # water that can be trapped at the index pointed to by the last moved pointer
            trappable = depth - temp

            # check if the trappable height is non-zero and non-negative
            if trappable > 0:
                # add the trappable height to the answer
                total += trappable

            # update the highest wall seen on the left size
            leftMax = max(leftMax, height[l])
            # update the highest wall seen on the right size
            rightMax = max(rightMax, height[r])

        # return the final answer
        return total

Complexity

  • Time: $O(n)$
    Since we process $n$ elements at most.

  • Space: $O(1)$
    Since we only use a few integer variables.


Mistakes I Made

I had to look this one up :(


And we are done.