How to get the water level in each bar of a histogram after raining?

Question

There is a well-known problem called Trapping Rain Water or something similar. In this problem, you have to find the maximum amount of water that can stay inside the given histogram after it rains. I have a different version of this problem. I want to find the water level in each bar of a histogram. For example, look at this picture:

The input of this histogram is this array: {0, 3, 0, 2, 0, 4, 0} The output (water level in each bar) is this array: {0, 0, 3, 1, 3, 0, 0}

Please note that each element of the output is the height of water in that bar. It does not matter where the water starts and ends.

How can I change the algorithm linked to the question to get this output?

The complexity should not increase. It should be O(n).

Answer 1

The original algorithm (with precalculation, which has O(𝑛) complexity) does not need much change. I will use as reference here the Python version of that algorithm, as contributed by Smitha Dinesh Semwal on GeeksForGeeks:

def findWater(arr, n): 
    left = [0]*n 
    right = [0]*n 
    water = 0
    left[0] = arr[0] 
    for i in range(1, n): 
        left[i] = max(left[i-1], arr[i]) 
    right[n-1] = arr[n-1] 
    for i in range(n-2, -1, -1): 
        right[i] = max(right[i + 1], arr[i]) 
    for i in range(0, n): 
        water += min(left[i], right[i]) - arr[i] 
    return water 

We see how the final loop adds vertical stacks of water to the final result. And as these vertical stacks are exactly what you want to get in an array (list) output, the change to that last loop is rather trivial:

    result = []
    for i in range(0, n): 
        result.append(min(left[i], right[i]) - arr[i])
    return result