Python Coding Test Course, Calculating the Amount of Water

One of the common problems that appears in coding tests is to find the “amount of water” that satisfies certain conditions. In this course, we will teach you how to write Python code and the importance of algorithmic thinking through this problem.

Problem Description

You need to calculate the amount of water in a single reservoir. An integer array height representing the height of the reservoir is given. The reservoir consists of blocks arranged horizontally, and the height of each block is expressed as height[i].

After it rains, you need to create a function that calculates the amount of water trapped between each block and returns the total amount of trapped water.

Input

  • height: A list of positive integers. (1 ≤ height.length ≤ 2 * 10^4, 0 ≤ height[i] ≤ 10^5)

Output

  • An integer representing the total amount of trapped water

Examples

Example 1

Input: height = [0,1,0,2,1,0,1,3,2,1,2,1]
Output: 6

Explanation: The amount of trapped water is 6.

Example 2

Input: height = [4,2,0,3,2,5]
Output: 9

Explanation: The amount of trapped water is 9.

Problem Solving Approach

To solve this problem, you need to understand two main points:

  1. To calculate the amount of trapped water at position i, you need to know the maximum heights on both sides of position i.
  2. The amount of water trapped at each position can be calculated as min(left maximum height, right maximum height) - current height.

To achieve this, we will create two arrays. One will store the maximum height from the left for each position, and the other will store the maximum height from the right.

Code Implementation

Now, let’s write the actual code based on this. 

def trap(height):
    if not height:
        return 0

    n = len(height)
    left_max = [0] * n
    right_max = [0] * n

    # Calculate left maximum height
    left_max[0] = height[0]
    for i in range(1, n):
        left_max[i] = max(left_max[i - 1], height[i])

    # Calculate right maximum height
    right_max[n - 1] = height[n - 1]
    for i in range(n - 2, -1, -1):
        right_max[i] = max(right_max[i + 1], height[i])

    # Calculate the amount of trapped water
    water_trapped = 0
    for i in range(n):
        water_trapped += min(left_max[i], right_max[i]) - height[i]

    return water_trapped

# Run examples
print(trap([0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]))  # Output: 6
print(trap([4, 2, 0, 3, 2, 5]))  # Output: 9

Code Explanation

1. trap(height): This function calculates the amount of water trapped.
2. It receives the height list as input and calculates the amount of water based on it.
3. Returns 0 if the list is empty.
4. Calculates the left maximum height at each index, and calculates the right maximum height.
5. Finally, it calculates the amount of trapped water using the stored maximum heights and the current height.

Time Complexity Analysis

The time complexity of this algorithm is O(n). It takes O(n) time to create the two arrays and calculate their maximum heights, and it takes another O(n) time to measure the trapped water.

Combining all steps results in a total of O(n).

Conclusion

In this course, we have learned in depth about the problem of finding the amount of water. This problem is a representative example of algorithm problem-solving and has various variations. We hope you will also practice similar problems to improve your coding skills.

Additional Practice Problems

It is recommended to practice the following variant problems:

  • Find the position where the most water gets trapped in the given height array.
  • Handle the case where it doesn’t rain, i.e., when the height at all positions is the same.
  • Analyze how it changes when there is water flow (i.e., when the water in each block can flow into adjacent blocks).