Накопитель дождевой воды

Источник: https://leetcode.com/problems/trapping-rain-water/

Даны nнеотрицательных чисел, которые кодируют неровную поверхность, состояющую из столбиков. Ширина каждого столбика = 1. Подсчитать, сколько дождевой воды можно накопить между столбиками.  

Пример 1:

Дано: height = [0,1,0,2,1,0,1,3,2,1,2,1]
Результат: 6
Объяснение: Неровная поверхность, образуемая столбиками (черные квадраты) закодирована массивом [0,1,0,2,1,0,1,3,2,1,2,1]. Вода (голубые столбики) задерживается во впадинах паверхности. И обозначена голубыми квадратами. Всего у нас 6 голубых квадратов.

Пример 2:

Дано: height = [4,2,0,3,2,5]
Результат: 9

Ограничения:

• n == height.length
• 1 <= n <= 2 * 104
• 0 <= height[i] <= 105

Идея решения

Мы можем использовать стек, чтобы учитывать столбики поверхности, которые окружены другими столбиками - и значит, они могут быть покрыты водой.

Сканируем массив и добавляем индекс элемента в стек, когда он меньше или равен элементу, соответствующему индексу в топе стека, это значит, что элемент топа стека создает стенку для текущего элемента. Если же текущий элемент больше элемента в топе стека, то текущий элемент является стенкой для топового элемента стека, и потому мы можем вытолкнуть его из стека, подсчитать собранную воду и добавить ее к ответу.

Алгоритм

• Используем стек.
• IСканируем массив:
  • Пока стек не пустой и $\text{height[current]}>\text{height[st.top()]}$
    • Выталкиваем элемент  стека и называем его $\text{top}$.
      
    • Считаем растсояние между текущим элементом массива и оставшимся топовым элементом стека: $\text{distance} = \text{current} - \text{st.top}() - 1$
    • Высота столбика собранной воды: $\text{bounded\_height} = \min(\text{height[current]}, \text{height[st.top()]}) - \text{height[top]}$
    • Аккумулируем ответ, подсчитывая количество голубых квадратов: $\text{ans} \mathrel{+}= \text{distance} \times \text{bounded\_height}$
  • Поместить текущий индекс в стек
  • Перейти к следующему элементу
    

Имплеменация

class Solution {
    public int trap(int[] height) {
        int ans = 0, current = 0;
        Stack<Integer> st = new Stack<Integer>();
        while (current < height.length) {
            while (!st.empty() && height[current] > height[st.peek()]) {
                int top = st.peek();
                st.pop();
                if (st.isEmpty()) // у топового элемента нет стенки слева
                    break;
                int distance = current - st.peek() - 1;
                int bounded_height = Math.min(height[current], height[st.peek()]) - height[top];
                ans += distance * bounded_height;
            }
            st.push(current++);
        }
        return ans;
    }
}

Оценка сложности

• Сложность по времени: $O(n)$.
  • Мы сканируем массив один раз, поэтому $O(n)$
• Сложность по памяти: $O(n)$. Стек может знаять не более $O(n)$ памяти в случае лестнице-подобной или плоской структуры поверхности, закодированной исходным массивом.