Compiler Explorer

Source code

pub fn compute_rain_collected(height: &[i64]) -> u64 {
    let mut height = height.into_iter();

std::iter::repeat(())
        // We reorder the sequence of elevations by taking values
        // from both ends of the terrain on a minimum first basis,
        // advancing the iterator that points to the smaller value.
        // This way we are guaranteed to have the maximum elevation
        // as the last spot.
        .scan((height.next(), height.next_back()), |state, _| {
            if let (Some(left), Some(right)) = *state {
                if left <= right {
                    *state = (height.next(), Some(right));
                    Some(left)
                } else {
                    *state = (Some(left), height.next_back());
                    Some(right)
                }
            } else {
                None
            }
        })
        // The newly ordered sequence traps the same amount of
        // water as the former one (yes, you got it, proof is left
        // to the reader). We can now apply the fold operation
        // described above for the "maximum is last" case.
        .fold(
            (
                i64::MIN, // keeps track of the stepsize of the stair
                0u64,     // keeps track of the water collected
            ),
            |acc, &x| {
                let stepsize = x.max(acc.0);
                (stepsize, acc.1 + (stepsize - x) as u64)
            },
        )
        .1 // we are only interested in the water collected
}