Trapping Rain Water: сколько воды удержат столбцы — от O(n²) до двух указателей

Коротко

Trapping Rain Water — одна из самых узнаваемых задач уровня Hard. Её любят за то, что она проверяет сразу всё: умеете ли вы найти локальную формулу, улучшить решение по шагам и доказать, почему жадный ход с двумя указателями корректен. Ниже — весь путь от первой идеи до оптимума и обобщение на матрицу.

Условие

Дан массив неотрицательных чисел height: это ширина-1 столбцы рельефа. После дождя между столбцами скапливается вода. Нужно посчитать, сколько единиц воды удержится.

height = [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]  →  6
height = [4, 2, 0, 3, 2, 5]                    →  9

На LeetCode это задача 42, длина массива до $2 \cdot 10^4$, высоты до $10^5$. Квадратичное решение там формально проходит, но на собеседовании его примут только как отправную точку.

Нарисуем второй пример. # — столбец, ~ — вода:

.....#
#~~~~#
#~~#~#
##~###
##~###
420325

Точнее по столбцам: над 2 стоит 2 единицы воды, над 0 — 4, над 3 — 1, над 2 — 2. Итого 9.

Интуиция: смотрим на один столбец

Главная ошибка — пытаться искать «ямы» целиком: где начинается, где кончается, какой у неё уровень. Ямы бывают вложенными, и такой код быстро превращается в кашу.

Проще думать о каждом столбце отдельно. Вода над столбцом i поднимется до уровня, который удержат стенки по бокам. Левая стенка — самый высокий столбец слева (включая сам i), правая — самый высокий справа. Вода переливается через более низкую из двух:

Мы включаем сам столбец в оба максимума, поэтому разность никогда не бывает отрицательной: если столбец выше соседей, над ним просто 0.

Вся задача сводится к тому, как быстро получить эти два максимума для каждого i.

Наивное решение: O(n²)

Прямо по формуле: для каждого столбца пробегаем влево и вправо.

def trap_naive(height):
    n = len(height)
    water = 0
    for i in range(n):
        left_max = max(height[: i + 1])
        right_max = max(height[i:])
        water += min(left_max, right_max) - height[i]
    return water

Для каждого из $n$ столбцов мы делаем $O(n)$ работы, итого $O(n^2)$. Срезы в Python ещё и копируют данные, так что по памяти это $O(n)$ на итерацию; в варианте с явными циклами было бы $O(1)$.

Это решение стоит проговорить вслух на собеседовании: оно показывает, что вы нашли правильную формулу, а дальше остаётся инженерная работа — убрать повторные вычисления.

Префиксные максимумы: O(n) времени и O(n) памяти

Повторная работа очевидна: максимум слева для i — это максимум слева для i − 1, сравненный с height[i]. Посчитаем оба массива заранее, один проход слева направо, другой справа налево. Это та же идея, что у префиксных сумм, только с max вместо +.

def trap_prefix(height):
    n = len(height)
    if n == 0:
        return 0
    left = [0] * n
    right = [0] * n
    left[0] = height[0]
    for i in range(1, n):
        left[i] = max(left[i - 1], height[i])
    right[n - 1] = height[n - 1]
    for i in range(n - 2, -1, -1):
        right[i] = max(right[i + 1], height[i])
    return sum(min(left[i], right[i]) - height[i] for i in range(n))

Для примера [4, 2, 0, 3, 2, 5]:

Сумма последней строки — 9. Три прохода по массиву, $O(n)$ времени и два вспомогательных массива, $O(n)$ памяти. Похожий приём с проходами в обе стороны используется в Product of Array Except Self.

Многие интервьюеры на этом и останавливаются. Но если спросят «а можно без лишней памяти?» — нужны два указателя.

Два указателя: O(n) времени и O(1) памяти

Идея

Посмотрим на формулу ещё раз: нам нужен не каждый из максимумов, а только меньший из них. Это и позволяет обойтись без массивов.

Поставим указатели left и right на края и будем хранить left_max — максимум на [0..left] и right_max — максимум на [right..n−1].

Допустим, height[left] < height[right]. Тогда:

• left_max — точный максимум слева для столбца left.
• Настоящий максимум справа для left не меньше height[right], потому что столбец right лежит справа от left.

Осталось показать, что left_max не больше этого настоящего максимума. Здесь помогает инвариант: один из указателей всегда стоит на самом высоком из уже просмотренных столбцов. Двигается всегда меньшая сторона, поэтому указатель на максимуме может сдвинуться, только если второй указатель стоит на столбце не ниже — и тогда максимум уже под ним.

Раз height[left] < height[right], максимум просмотренного не может стоять под левым указателем, значит, он под правым: left_max ≤ height[right] ≤ max(height[right..n−1]). Поэтому вода над left равна ровно left_max − height[left]. Точный правый максимум знать не нужно — достаточно знать, что он не меньше.

Симметрично, если правый столбец не выше левого, обрабатываем правую сторону.

Код

def trap(height):
    left, right = 0, len(height) - 1
    left_max = right_max = 0
    water = 0
    while left < right:
        if height[left] < height[right]:
            left_max = max(left_max, height[left])
            water += left_max - height[left]
            left += 1
        else:
            right_max = max(right_max, height[right])
            water += right_max - height[right]
            right -= 1
    return water

То же на JavaScript, если собеседование идёт на нём:

function trap(height) {
  let left = 0
  let right = height.length - 1
  let leftMax = 0
  let rightMax = 0
  let water = 0
  while (left < right) {
    if (height[left] < height[right]) {
      leftMax = Math.max(leftMax, height[left])
      water += leftMax - height[left]
      left++
    } else {
      rightMax = Math.max(rightMax, height[right])
      water += rightMax - height[right]
      right--
    }
  }
  return water
}

Пустой массив и массив из одного столбца обрабатываются сами: цикл не выполнится ни разу, ответ 0. Последний столбец, на котором встречаются указатели, — глобальный максимум, вода над ним всегда 0, поэтому условие left < right, а не <=.

Разбор на примере

height = [4, 2, 0, 3, 2, 5]. Правый край 5 выше всех, поэтому всё время двигается левый указатель:

После пятого шага left == right == 5, цикл заканчивается. Ответ 9.

Обратите внимание: мы ни разу не узнали, что максимум справа для столбца 1 — это 5. Нам хватило того, что он не меньше 4.

Если вы решали Container With Most Water, схема покажется знакомой: там тоже двигают меньшую сторону, потому что именно она ограничивает ответ. Общая техника разобрана в статье про два указателя.

Монотонный стек: считаем воду слоями

Есть совсем другой взгляд на ту же картинку. Вместо того чтобы считать воду вертикальными столбиками, можно считать её горизонтальными слоями: каждый раз, когда мы находим «дно», ограниченное стенками слева и справа, заливаем прямоугольник над ним.

Для этого нужен монотонный стек индексов с невозрастающими высотами. Когда приходит столбец выше вершины, вершина становится дном ямы:

• снимаем со стека дно bottom;
• если стек опустел, левой стенки нет, вода утекает;
• иначе левая стенка — новая вершина left, правая — текущий i;
• ширина слоя i − left − 1, высота min(height[left], height[i]) − height[bottom].

def trap_stack(height):
    stack = []  # индексы, высоты по убыванию
    water = 0
    for i, h in enumerate(height):
        while stack and height[stack[-1]] < h:
            bottom = stack.pop()
            if not stack:
                break
            left = stack[-1]
            width = i - left - 1
            depth = min(height[left], h) - height[bottom]
            water += width * depth
        stack.append(i)
    return water

На [4, 2, 0, 3, 2, 5] стек заливает четыре прямоугольника:

Каждый индекс кладётся в стек и снимается не больше одного раза, поэтому время $O(n)$, память $O(n)$ в худшем случае (убывающий массив).

Зачем это знать, если два указателя проще и экономнее? Две причины. Во-первых, интервьюер может попросить «другое решение» — стек хорошо показывает кругозор. Во-вторых, послойный подсчёт обобщается на задачи, где нужно именно «дно и две стенки»: Largest Rectangle in Histogram и родственные ей. А если хочется размяться на более простой задаче со стеком, начните с Daily Temperatures.

Сравнение подходов

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

Вариант 2D: Trapping Rain Water II

В LeetCode 407 рельеф задан матрицей heightMap, и вода может растекаться в четыре стороны. Формула «минимум из двух максимумов» здесь не работает: у клетки нет «лева» и «права», вода утекает по самому низкому пути к краю.

Идея — заливать снаружи внутрь, как вода в Дейкстре:

1. Все граничные клетки кладём в минимальную кучу — это стенка бассейна, через неё вода уходит наружу.
2. Достаём клетку с минимальной высотой level. Это самое слабое место текущей стенки: ни одна внутренняя соседняя клетка не может удержать воду выше level.
3. Для каждого непосещённого соседа: если он ниже level, над ним встанет level − height воды. В кучу он идёт с высотой max(level, height) — залитая клетка становится частью стенки на уровне воды.

import heapq


def trap_rain_water_2d(grid):
    if not grid or not grid[0]:
        return 0
    m, n = len(grid), len(grid[0])
    seen = [[False] * n for _ in range(m)]
    heap = []
    for r in range(m):
        for c in range(n):
            if r in (0, m - 1) or c in (0, n - 1):
                heapq.heappush(heap, (grid[r][c], r, c))
                seen[r][c] = True
    water = 0
    while heap:
        level, r, c = heapq.heappop(heap)
        for nr, nc in ((r + 1, c), (r - 1, c), (r, c + 1), (r, c - 1)):
            if 0 <= nr < m and 0 <= nc < n and not seen[nr][nc]:
                seen[nr][nc] = True
                water += max(0, level - grid[nr][nc])
                heapq.heappush(heap, (max(level, grid[nr][nc]), nr, nc))
    return water

Для матрицы

1 4 3 1 3 2
3 2 1 3 2 4
2 3 3 2 3 1

ответ 4: вода встаёт в трёх клетках средней строки — 1 над левой двойкой, 2 над единицей и 1 над правой двойкой. Каждая клетка попадает в кучу один раз, поэтому сложность $O(mn \log(mn))$ по времени и $O(mn)$ по памяти. Подробнее про кучу — в статьях о бинарной куче и алгоритме Дейкстры: здесь, по сути, та же жадность «обработать самую низкую точку фронта».

Кстати, одномерную задачу тоже можно решать этим способом: на краях стоят два граничных столбца, и «куча» из двух элементов — это ровно два указателя. Хороший ответ на вопрос «а как связаны ваши решения».

Типичные ошибки

• Искать ямы целиком. Код с поиском левой и правой стенки каждой ямы почти всегда ломается на вложенных ямах вроде [5, 1, 3, 1, 5]. Думайте о столбце, а не о яме.
• Не включать сам столбец в максимум. Если left_max считать только по элементам строго левее, разность бывает отрицательной, и придётся писать max(0, …). Работает, но лишняя ветка — лишний шанс ошибиться.
• Двигать не ту сторону. В двух указателях двигается сторона с меньшей высотой. Если перепутать, рассуждение о том, что противоположная сторона «не ниже», перестаёт выполняться, и ответ завышается.
• Забыть break в стековом решении. Если после снятия дна стек пуст, левой стенки нет. Без проверки stack[-1] упадёт с ошибкой или, хуже, возьмёт мусор.
• Путать с Container With Most Water. Там выбирают две стенки и считают одну площадь, столбцы между ними не мешают. Здесь вода считается над каждым столбцом, и промежуточные столбцы вытесняют воду.
• В 2D добавлять в кучу исходную высоту. Если залитая клетка уходит в кучу с grid[r][c], а не с max(level, grid[r][c]), соседние клетки будут считать, что стенка ниже, чем есть, и воды получится меньше.

Trapping Rain Water на собеседовании

Задача есть в популярных подборках вроде NeetCode 150 и считается одной из классических Hard-задач для алгоритмической секции. Удобный план ответа на 30–40 минут:

1. Уточнить условие: высоты неотрицательные, ширина столбцов 1, вода с краёв утекает.
2. Нарисовать пример и вывести формулу min(leftMax, rightMax) − h. Это самый важный момент — без него дальше двигаться некуда.
3. Озвучить наивное $O(n^2)$ и сразу предложить префиксные максимумы.
4. Перейти к двум указателям и словами объяснить, почему двигаем меньшую сторону. Именно это объяснение отличает понимание от заученного кода.
5. Прогнать код на маленьком примере и на краевых: пустой массив, монотонный массив, все столбцы одной высоты.

Частые follow-up вопросы:

• «Можно без доппамяти?» — два указателя.
• «А другим способом?» — монотонный стек, послойный подсчёт.
• «А в двумерном случае?» — куча от границы, сложность $O(mn \log(mn))$.
• «А если массив приходит потоком?» — ответ зависит от будущих столбцов, поэтому в общем случае нужно хранить незакрытую часть рельефа; удобно это делает тот же монотонный стек, который заливает воду, как только появляется правая стенка.

Итог

Trapping Rain Water учит важной привычке: сначала найти формулу для одного элемента, а потом убирать из неё повторные вычисления. Формула min(leftMax, rightMax) − h даёт наивное решение, префиксные массивы делают его линейным, а наблюдение «важен только меньший максимум» сводит память к $O(1)$. Монотонный стек и куча показывают ту же задачу с других сторон и пригодятся в соседних задачах.

Перед собеседованием напишите два указателя с нуля без подсказок и проверьте себя на [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1] → 6 и [4, 2, 0, 3, 2, 5] → 9.