1 def print_terrain(heights):
 2     if not heights:
 3         return
 4 
 5     max_height = max(heights)
 6 
 7     # Iterate from the highest point down to 1
 8     for level in range(max_height, 0, -1):
 9         row = ""
10         for h in heights:
11             if h >= level:
12                 row += "+"
13             else:
14                 row += " "
15         print(row)
16 
17     # Print the base layer (level 0)
18     print("+" * len(heights) + " <--- base layer")
19 
20 # The data
21 data = [5, 4, 3, 2, 1, 3, 4, 0, 3, 4]
22 print_terrain(data)

 1 def dump_water(heights, water_amount, column):
 2     # Use a separate list to track water so we can render '+' and 'W' separately
 3     water = [0] * len(heights)
 4     
 5     for _ in range(water_amount):
 6         curr = column
 7         
 8         # 1. Try to move Left
 9         best_left = curr
10         for i in range(curr - 1, -1, -1):
11             if heights[i] + water[i] < heights[best_left] + water[best_left]:
12                 best_left = i
13             elif heights[i] + water[i] > heights[best_left] + water[best_left]:
14                 break # Hit a wall
15         
16         if best_left < curr:
17             water[best_left] += 1
18             continue # Water unit settled to the left
19 
20         # 2. Try to move Right
21         best_right = curr
22         for i in range(curr + 1, len(heights)):
23             if heights[i] + water[i] < heights[best_right] + water[best_right]:
24                 best_right = i
25             elif heights[i] + water[i] > heights[best_right] + water[best_right]:
26                 break # Hit a wall
27         
28         # Regardless of if it moved right or stayed at 'curr', increment
29         water[best_right] += 1
30 
31     return water
32 
33 def render_water_terrain(heights, water):
34     max_h = max(h + w for h, w in zip(heights, water))
35     
36     for level in range(max_h, 0, -1):
37         row = ""
38         for i in range(len(heights)):
39             if level <= heights[i]:
40                 row += "+"
41             elif level <= heights[i] + water[i]:
42                 row += "W"
43             else:
44                 row += " "
45         print(row)
46     print("+" * len(heights) + " <--- base layer")
47 
48 # Execution
49 terrain = [5, 4, 3, 2, 1, 3, 4, 0, 3, 4]
50 water_dist = dump_water(terrain, water_amount=8, column=1)
51 render_water_terrain(terrain, water_dist)