Дан ряд целых чисел $a_i \geqslant 0,$ $i=\overline{1,n}$. Числа символизируют высоту стен. Сверху идёт дождь, и вода скапливается в ячейках между стенами. В каждом столбце скапливается $w_i$ воды. Нужно найти количество ячеек с водой: $$W=\sum_{i=1}^{n} w_i.$$

Пример: $$a=(2, 4, 1, 3, 2, 5, 2, 1, 3, 1).$$

⬛⬛⬛⬛⬛⬜⬛⬛⬛⬛

⬛⬜🟦🟦🟦⬜⬛⬛⬛⬛

⬛⬜🟦⬜🟦⬜🟦🟦⬜⬛

⬜⬜🟦⬜⬜⬜⬜🟦⬜⬛

⬜⬜⬜⬜⬜⬜⬜⬜⬜⬜

В столбцах скопилось воды $w=(0,0,3,1,2,0,1,2,0,0),$ в сумме $W = 9.$

Решение

Разобьём исходный массив на два по максимальному значению.

⬛⬛⬛⬛⬛🟥⬛⬛⬛⬛

⬛⬜🟦🟦🟦🟥⬛⬛⬛⬛

⬛⬜🟦⬜🟦🟥🟦🟦⬜⬛

⬜⬜🟦⬜⬜🟥⬜🟦⬜⬛

⬜⬜⬜⬜⬜🟥⬜⬜⬜⬜

 →⬛⬛⬛⬛⬛⬛⬛⬛ ←

Индекс этого значения: $k=\underset{i}{\mathrm{argmax}}\, a_i,$ $k=6$. Подмассивы размерами $n_1=5$ и $n_2=4$ для этого примера:

$$a^{(L)}=(2,4,1,3,2); a^{(R)}=(2,1,3,1).$$

Тогда в левой части максимальный уровень воды — максимум от левого края до текущего элемента. В правой части максимальный уровень воды — максимум от текущего элемента до правого края:

$$l_{i}^{(L)}=\max_{1\leqslant j\leqslant i} a_j^{(L)};l_{i}^{(R)}=\max_{i\leqslant j\leqslant n_2} a_j^{(R)}.$$

В примере это:

$$l^{(L)}=(2,4,4,4,4); l^{(R)}=(3,3,3,1).$$

Считаем разности между максимальным уровнем воды и высотой стены, суммируем:

$$W=\sum_{i=1}^{n_1} \left(l_{i}^{(L)}-a_i^{(L)}\right)+\sum_{i=1}^{n_2} \left(l_{i}^{(R)}-a_i^{(R)}\right).$$

Количество воды в каждом столбце:

$$w^{(L)}=l^{(L)}-a^{(L)}=(0,0,3,1,2); w^{(R)}=l^{(R)}-a^{(R)}=(1,2,0,0).$$

Сложность решения по времени — $O(n)$, по памяти — $O(1)$.

C++20 и Range-v3

Индекс максимального элемента — расстояние от начала вектора до максимального элемента:

auto indexMax{ ranges::distance(walls.begin(), ranges::max_element(walls)) };

Поиск максимумов $l^{(L)}$ и $l^{(R)}$ очень похож, поэтому хватит одной универсальной лямбда-функции. Состояние реализовано через захват целого level, а чтобы его можно было изменять, лямбда-функция объявлена как mutable.

auto difference{ [level{ 0 }](auto wall) mutable noexcept {

    if (wall > level)

    {

        level = wall;

    }

    return level - wall;

} };

Количество воды в левой части — сумма разностей между высотой воды и высотой стены:

ranges::accumulate(

    walls

    | ranges::views::take(indexMax)

    | ranges::views::transform(difference), 0)

level        [2] [4] [4] [4] [4]

wall          2   4   1   3   2

              ↓   ↓   ↓   ↓   ↓

level - wall  0 + 0 + 3 + 1 + 2 = 6

Количество воды в правой части считается аналогично, но входные данные переворачиваются.

ranges::accumulate(

    walls

    | ranges::views::drop(indexMax + 1)

    | ranges::views::reverse

    | ranges::views::transform(difference), 0)

level        [1] [3] [3] [3]

wall          1   3   1   2

              ↓   ↓   ↓   ↓

level - wall  0 + 0 + 2 + 1 = 3

С помощью алгоритмов на диапазонах можно проверить, корректны ли входные данные.

auto isNegative{ std::bind(std::less{}, std::placeholders::_1, 0) };

if (ranges::any_of(walls, isNegative))

{

    throw std::invalid_argument{ "< 0" };

}

Полный код.

int countWaterCells(const std::vector<int>& walls) noexcept(false)

{

    auto isNegative{ std::bind(std::less{}, std::placeholders::_1, 0) };

    if (ranges::any_of(walls, isNegative))

    {

        throw std::invalid_argument{ "< 0" };

    }

 

    auto difference{ [level{ 0 }](auto wall) mutable noexcept {

        if (wall > level)

        {

            level = wall;

        }

        return level - wall;

    } };

 

    auto indexMax{ ranges::distance(walls.begin(), ranges::max_element(walls)) };

    return

        ranges::accumulate(

            walls

            | ranges::views::take(indexMax)

            | ranges::views::transform(difference), 0)

        + ranges::accumulate(

            walls

            | ranges::views::drop(indexMax + 1)

            | ranges::views::reverse

            | ranges::views::transform(difference), 0);

}

Приложение: исходный код.