Given n non-negative integers representing an elevation map where the width of each bar is 1, compute how much water it is able to trap after raining.

rainy_valleys.py

# -*- coding: utf-8 -*-
"""
Created on Mon Mar  9 22:51:09 2020
Given n non-negative integers representing an elevation map where the width
of each bar is 1, compute how much water it is able to trap after raining.

@author: Hemerson Tacon
"""


def rainy_valleys(elevation_map):

    if len(elevation_map) < 3:
        return 0

    water = 0
    
    local_maxima = []
    global_maximum = -1
    
    for idx in range(len(elevation_map)):
        if discrete_local_maximum(idx, elevation_map):
            local_maxima.append(idx)
            
            if global_maximum == -1 or elevation_map[idx] > elevation_map[global_maximum]:
                global_maximum = idx
                
    left_portion = local_maxima[:local_maxima.index(global_maximum)]
    last_max_idx_in_map = global_maximum
    
    while len(left_portion) > 0:
        hills = [elevation_map[idx] for idx in left_portion]
        curr_max_hill_idx = hills.index(max(hills))
        curr_max_idx_in_map = left_portion[curr_max_hill_idx]
        
        right_hill = elevation_map[last_max_idx_in_map]
        left_hill = elevation_map[curr_max_idx_in_map]
        
        valley_length = last_max_idx_in_map - curr_max_idx_in_map - 1
        valley_height = min(left_hill, right_hill)
        rocks_amount = sum([min(x, valley_height) for x in elevation_map[curr_max_idx_in_map + 1:last_max_idx_in_map]])
        water += valley_height * valley_length - rocks_amount
        
        left_portion = left_portion[:curr_max_hill_idx]
        last_max_idx_in_map = curr_max_idx_in_map
        
    right_portion = local_maxima[local_maxima.index(global_maximum) + 1:]
    last_max_idx_in_map = global_maximum
    
    while len(right_portion) > 0:
        hills = [elevation_map[idx] for idx in right_portion]
        curr_max_hill_idx = hills.index(max(hills))
        curr_max_idx_in_map = right_portion[curr_max_hill_idx]
        
        left_hill = elevation_map[last_max_idx_in_map]
        right_hill = elevation_map[curr_max_idx_in_map]
        
        valley_length = curr_max_idx_in_map - last_max_idx_in_map - 1
        valley_height = min(left_hill, right_hill)
        rocks_amount = sum([min(x, valley_height) for x in elevation_map[last_max_idx_in_map + 1:curr_max_idx_in_map]])
        water += valley_height * valley_length - rocks_amount
    
        right_portion = right_portion[curr_max_hill_idx + 1:]
        last_max_idx_in_map = curr_max_idx_in_map

    return water


def discrete_local_maximum(idx, values):

    if(idx == 0):
        return check_right(idx, values)

    elif(idx == len(values) - 1):
        return check_left(idx, values)

    else:
        return check_right(idx, values) and check_left(idx, values)


def check_right(idx, values):

    adjacent_idx = idx + 1

    while values[idx] == values[adjacent_idx]:
        adjacent_idx += 1
        if adjacent_idx == len(values):
            return True

    return values[idx] > values[adjacent_idx]


def check_left(idx, values):

    adjacent_idx = idx - 1

    while values[idx] == values[adjacent_idx]:
        adjacent_idx -= 1
        if adjacent_idx == -1:
            return True

    return values[idx] > values[adjacent_idx]


if __name__ == "__main__":

    x = [6, 4, 1, 2, 3, 0, 1, 2, 6, 0, 3, 3, 3, 0, 5]
    print(f"Input: {x}")
    print(f"Output: {rainy_valleys(x)}")