House Robber (грабитель домов): от рекурсии к O(1) памяти и вариант по кругу

Коротко

• Задача: из ряда домов выбрать несоседние так, чтобы сумма денег была максимальной.
• Главная идея: на каждый дом есть два варианта. Взять его и прибавить лучший результат до i - 2 или пропустить и оставить лучший результат до i - 1.
• Сложность: $O(n)$ по времени и $O(1)$ по памяти, потому что нужны только два предыдущих значения.
• Жадность не работает: на [2, 3, 2] самый дорогой дом даёт 3, а правильный ответ 4.
• Дома по кругу (House Robber II): первый и последний дом теперь соседи, поэтому считаем линейную версию дважды: без первого дома и без последнего.

Условие

Дома стоят в ряд, в каждом лежит nums[i] денег. Если ограбить два соседних дома, сработает сигнализация. Какую максимальную сумму можно унести?

nums = [2, 7, 9, 3, 1]  →  12   (дома 0, 2, 4: 2 + 9 + 1)
nums = [1, 2, 3, 1]     →  4    (дома 0, 2: 1 + 3)

На LeetCode это задача 198. Решить её можно и у нас, она есть в блоке «Задачи по теме». Ограничения там небольшие: до 100 домов и суммы до 400, но на собеседовании от вас всё равно будут ждать линейное решение.

Почему жадный выбор не работает

Первая мысль, которая приходит в голову: брать самое дорогое, вычёркивать соседей, повторять. Запишем это и прогоним на примерах из условия.

def greedy(nums):
    order = sorted(range(len(nums)), key=lambda i: nums[i], reverse=True)
    taken = set()
    for i in order:
        if i - 1 not in taken and i + 1 not in taken:
            taken.add(i)
    return sum(nums[i] for i in taken)


print(greedy([2, 7, 9, 3, 1]))  # 12
print(greedy([1, 2, 3, 1]))  # 4
print(greedy([2, 3, 2]))  # 3

Оба примера из условия проходят, и в этом главная ловушка: кажется, что решение верное. Третий тест его ломает. Жадный алгоритм берёт дом за 3 и теряет двух соседей, хотя вместе они стоят 4.

Вторая популярная идея: дома нельзя брать подряд, значит, ответ это сумма либо чётных, либо нечётных позиций. Эта версия тоже проходит оба примера и падает на другом тесте. На [2, 1, 1, 2] выгодно пропустить сразу два дома подряд и взять крайние.

Причина одна: цена дома сама по себе ничего не говорит о выгоде. Взяв дом, мы платим за него двумя соседями, а эта плата зависит от того, что выбрано вокруг. Безопасного локального выбора, какой нужен жадным алгоритмам, здесь нет: решение по дому i меняет то, что выгодно делать с домом i + 1.

Первая жадная стратегия при этом не безнадёжна: она всегда набирает не меньше половины оптимума. Каждый дом из оптимального ответа либо взят ею, либо стоит рядом с взятым домом не дешевле себя, а соседей у дома максимум два. На [100, 101, 100] получается 101 против 200, так что граница почти достигается. Для задачи, где нужен точный максимум, этого мало.

Контрпример полезно назвать на собеседовании вслух: так видно, что вы проверили идею, прежде чем писать код.

Шаг 1. Рекурсия: формулируем выбор

Посмотрим на последний дом с индексом i. Вариантов ровно два:

• пропускаем его: ответ тот же, что для домов 0..i-1;
• берём его: дом i - 1 закрыт, получаем nums[i] плюс ответ для домов 0..i-2.

Других вариантов нет, поэтому перебор полный. Это и есть оптимальная подструктура: лучший ответ на префиксе собирается из лучших ответов на более коротких префиксах. Сколько именно домов пропустили подряд, отдельно учитывать не надо: best(i - 1) уже может заканчиваться пропуском.

Формула переводится в код один в один:

def rob_recursive(nums):
    def best(i):
        if i < 0:
            return 0
        return max(best(i - 1), best(i - 2) + nums[i])

    return best(len(nums) - 1)


print(rob_recursive([2, 7, 9, 3, 1]))  # 12

Код верный, но медленный. best(4) вызывает best(3) и best(2), а best(3) снова вызывает best(2). Такое дерево растёт как числа Фибоначчи, то есть как $O(\varphi^n)$, где $\varphi \approx 1.618$:

До ограничения в 100 домов такому решению не дожить. При этом различных аргументов у best всего около $n$: всё остальное время функция пересчитывает то, что уже считала.

Шаг 2. Мемоизация

Раз подзадачи повторяются, запомним ответы. Подробно приём разобран в статье про мемоизацию, здесь хватит пары строк:

from functools import lru_cache


def rob_memo(nums):
    @lru_cache(maxsize=None)
    def best(i):
        if i < 0:
            return 0
        return max(best(i - 1), best(i - 2) + nums[i])

    return best(len(nums) - 1)


print(rob_memo([2, 7, 9, 3, 1]))  # 12

Декоратор стоит на вложенной функции намеренно. В кэш попадает только i, а список nums для lru_cache не годится, потому что он нехешируемый. Заодно кэш живёт один вызов и не смешивает ответы для разных массивов.

Время теперь $O(n)$, память тоже $O(n)$: кэш плюс стек вызовов.

Со стеком есть нюанс. Глубина рекурсии равна числу домов, а лимит Python по умолчанию равен 1000. На Python 3.9 в моём прогоне версия с lru_cache работала при 490 домах и падала с RecursionError при 500: обёртка кэша расходует стек наравне с самой функцией. Версия с обычным словарём дотянула до 990. В других версиях Python пороги могут отличаться. Для 100 домов это не проблема, но упомянуть об ограничении стоит.

Шаг 3. Таблица

Избавимся от стека: будем считать ответы от коротких префиксов к длинным. Пусть dp[i] это лучший результат на первых i домах. Сдвиг на единицу убирает возню с отрицательными индексами: dp[0] = 0, домов нет и денег нет.

def rob_table(nums):
    n = len(nums)
    dp = [0] * (n + 1)
    dp[1] = nums[0]
    for i in range(2, n + 1):
        dp[i] = max(dp[i - 1], dp[i - 2] + nums[i - 1])
    return dp[n]


print(rob_table([2, 7, 9, 3, 1]))  # 12

Разбор на [2, 7, 9, 3, 1]:

На строке i = 4 алгоритм делает то, чего не умеет жадность: дом за 3 не окупается, и результат просто переносится дальше.

Какие именно дома взяли

Частый дополнительный вопрос. Таблица для этого и нужна: идём с конца и смотрим, откуда пришло значение. Если dp[i] == dp[i - 1], дом можно пропустить, иначе он взят, и мы перепрыгиваем через соседа.

def rob_with_houses(nums):
    n = len(nums)
    dp = [0] * (n + 1)
    dp[1] = nums[0]
    for i in range(2, n + 1):
        dp[i] = max(dp[i - 1], dp[i - 2] + nums[i - 1])

    houses = []
    i = n
    while i > 0:
        if dp[i] == dp[i - 1]:
            i -= 1
        else:
            houses.append(i - 1)
            i -= 2
    return dp[n], houses[::-1]


print(rob_with_houses([2, 7, 9, 3, 1]))  # (12, [0, 2, 4])
print(rob_with_houses([2, 1, 1, 2]))  # (4, [0, 3])

Версия с $O(1)$ памяти из следующего шага ответить на этот вопрос не сможет: она уже забыла промежуточные значения. Если на собеседовании просят список домов, оставляйте таблицу.

Шаг 4. O(1) памяти

В формуле участвуют только dp[i - 1] и dp[i - 2], остальная таблица не нужна. Заменим её двумя переменными: curr это лучший результат с учётом текущего дома, prev это результат на шаг раньше.

def rob(nums):
    prev, curr = 0, 0
    for x in nums:
        prev, curr = curr, max(curr, prev + x)
    return curr


print(rob([2, 7, 9, 3, 1]))  # 12
print(rob([1, 2, 3, 1]))  # 4
print(rob([2, 3, 2]))  # 4
print(rob([2, 1, 1, 2]))  # 4

Те же значения, что в таблице, только хранятся лишь последние два:

Обе переменные обновляются одной строкой, и это важно: правая часть считается целиком до присваивания. Если развести их на две строки, prev = curr затрёт старое значение, и curr начнёт прибавлять к себе каждый дом. На [2, 7, 9, 3, 1] такой код вернёт 22, то есть просто сумму массива.

В JavaScript для этого нужна временная переменная:

function rob(nums) {
  let prev = 0
  let curr = 0
  for (const x of nums) {
    const next = Math.max(curr, prev + x)
    prev = curr
    curr = next
  }
  return curr
}

console.log(rob([2, 7, 9, 3, 1])) // 12

То же самое через два состояния

Можно мыслить иначе и хранить не два соседних шага, а два исхода для текущего дома: take (дом взят) и skip (дом пропущен).

def rob_states(nums):
    take, skip = 0, 0
    for x in nums:
        take, skip = skip + x, max(take, skip)
    return max(take, skip)


print(rob_states([2, 7, 9, 3, 1]))  # 12

Решение то же самое, но эта форма лучше обобщается. Именно она переносится на деревья, как будет видно ниже.

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

Тот же путь проходит Climbing Stairs, а общая схема разобрана в статье про динамическое программирование для начинающих. Если вы знаете задачу о рюкзаке, выбор «взять или пропустить» уже знаком. Разница в том, что здесь вместо вместимости ограничение на соседей.

House Robber II: дома по кругу

Типовое продолжение (LeetCode 213): дома стоят по кругу, так что первый и последний тоже соседи.

Линейное решение теперь ошибается. На [2, 7, 9, 3, 1] оно отвечает 12 и берёт дома 0, 2 и 4, но 0 и 4 теперь рядом.

Цикл мешает, потому что у задачи нет начала: непонятно, с какого дома стартовать. Разорвать его помогает простое наблюдение: первый и последний дом не могут быть взяты одновременно. Значит, хотя бы один из них в ответе отсутствует, и любой допустимый набор домов лежит целиком в одном из двух линейных отрезков:

• nums[0..n-2], последний дом выброшен;
• nums[1..n-1], первый дом выброшен.

В каждом отрезке круга больше нет, и работает обычное решение. Берём максимум из двух. Набор домов, где нет ни первого, ни последнего, попадает в оба отрезка, но на максимум это не влияет.

def rob_range(nums, lo, hi):
    prev, curr = 0, 0
    for i in range(lo, hi):
        prev, curr = curr, max(curr, prev + nums[i])
    return curr


def rob_circle(nums):
    n = len(nums)
    if n == 1:
        return nums[0]
    return max(rob_range(nums, 0, n - 1), rob_range(nums, 1, n))


print(rob_circle([2, 3, 2]))  # 3
print(rob_circle([1, 2, 3, 1]))  # 4
print(rob_circle([2, 7, 9, 3, 1]))  # 11
print(rob_circle([5]))  # 5

Разбор на [2, 7, 9, 3, 1]: без последнего дома [2, 7, 9, 3] даёт 11, без первого [7, 9, 3, 1] даёт 10. Ответ 11.

Время $O(n)$, два прохода. Память $O(1)$, но только потому, что мы передаём границы. Вариант max(rob(nums[1:]), rob(nums[:-1])) короче, но срез копирует массив, и формально это уже $O(n)$ памяти. На собеседовании можно написать короткий вариант, если вы сами об этом скажете.

Три вещи, на которых здесь ошибаются:

Один дом. Без проверки n == 1 оба отрезка на [5] окажутся пустыми, и функция вернёт 0 вместо 5. Его легко пропустить: примеры из условия этот случай обычно не покрывают.

Попытка поправить линейный ответ. Идея «посчитаем как обычно и, если взяты оба крайних дома, выбросим более дешёвый» не работает. На [5, 1, 1, 5] линейный ответ равен 10 (оба крайних), после выбрасывания остаётся 5, а правильный ответ 6: например, дома 0 и 2. Выбросив дом, надо заново решать задачу для остальных.

Перебор всех точек разрыва. Запускать линейное решение для каждого стартового дома незачем, это $O(n^2)$. Достаточно двух запусков.

Есть и равносильная формулировка: зафиксировать судьбу дома 0. Либо берём его, тогда дома 1 и n - 1 закрыты, либо не берём, тогда круг уже разорван:

def rob_circle_fixed(nums):
    n = len(nums)
    return max(nums[0] + rob_range(nums, 2, n - 1), rob_range(nums, 1, n))


print(rob_circle_fixed([2, 7, 9, 3, 1]))  # 11
print(rob_circle_fixed([5]))  # 5

Здесь n == 1 отдельно проверять не нужно. Сам приём пригодится и в других задачах на цикл: зафиксировали состояние одного элемента, и цикл превратился в линию.

Что ещё могут спросить

Задачу удобно развивать, поэтому за ней часто идут уточняющие вопросы. Вот несколько вероятных.

Дома в дереве (House Robber III)

LeetCode 337: дома образуют двоичное дерево, нельзя брать вершину вместе с её родителем. Здесь пригодится форма с двумя состояниями: каждая вершина возвращает пару (взяли, не взяли).

class TreeNode:
    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


def rob_tree(root):
    def go(node):
        if node is None:
            return 0, 0
        left_take, left_skip = go(node.left)
        right_take, right_skip = go(node.right)
        take = node.val + left_skip + right_skip
        skip = max(left_take, left_skip) + max(right_take, right_skip)
        return take, skip

    return max(go(root))


root = TreeNode(3, TreeNode(2, None, TreeNode(3)), TreeNode(3, None, TreeNode(1)))
print(rob_tree(root))  # 7

Это частный случай динамики на деревьях: ряд домов и есть дерево, вытянутое в линию.

Delete and Earn: задача в маскировке

LeetCode 740: выбираете число x, получаете x очков, и все x - 1 и x + 1 удаляются. Сложите очки по значениям, и получится House Robber, где «домами» служат сами значения: взяли x, значит, соседние значения закрыты.

def delete_and_earn(nums):
    points = [0] * (max(nums) + 1)
    for x in nums:
        points[x] += x
    return rob(points)


print(delete_and_earn([3, 4, 2]))  # 6
print(delete_and_earn([2, 2, 3, 3, 3, 4]))  # 9

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

Минимизировать самую дорогую кражу (House Robber IV)

LeetCode 2560: нужно ограбить не меньше k домов так, чтобы самый дорогой из них был как можно дешевле. Это уже бинарный поиск по ответу: для порога cap проверяем, хватит ли домов не дороже него.

def min_capability(nums, k):
    def enough(cap):
        count, i = 0, 0
        while i < len(nums):
            if nums[i] <= cap:
                count += 1
                i += 2
            else:
                i += 1
        return count >= k

    lo, hi = min(nums), max(nums)
    while lo < hi:
        mid = (lo + hi) // 2
        if enough(mid):
            hi = mid
        else:
            lo = mid + 1
    return lo


print(min_capability([2, 3, 5, 9], 2))  # 5
print(min_capability([2, 7, 9, 3, 1], 2))  # 2

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

Другие изменения условия

Про отрицательные числа: rob([-1, -2]) вернёт 0, потому что пустой набор считается допустимым. Если условие его запрещает, нужна проверка из последней строки таблицы.

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

• Неверное определение dp[i]. Если считать, что дом i взят обязательно, то переход смотрит на dp[i-2] и dp[i-3], а ответом становится max(dp), а не последний элемент. Так тоже можно решить, но обычно ошибаются, смешивая эти определения: формула от одного, ответ от другого. Определение стоит записать словами до кода.
• Отрицательные индексы без сдвига. В таблице без сдвига на первых шагах появляются dp[-1] и dp[-2]. Python не упадёт, а прочитает конец массива: там пока нули, и результат окажется верным случайно. В JavaScript та же строка даст undefined, и весь массив превратится в NaN.
• Пустой массив. dp[1] = nums[0] упадёт на []. В условии задачи n >= 1, но на собеседовании лучше спросить об этом. Версия с двумя переменными вернёт 0 сама.
• Раздельное обновление переменных. Ошибка из шага 4: получается сумма всего массива.
• Забытый n == 1 в круговой версии. Возвращается 0 вместо единственного дома.
• Проверка только на примерах из условия. Обе жадные версии проходят оба примера. Добавляйте тесты вроде [2, 3, 2] и [2, 1, 1, 2].

Где это встречается на собеседованиях

House Robber и House Robber II входят в популярную подборку Blind 75, обе задачи на LeetCode средней сложности. Первая показывает, понимает ли кандидат динамику, вторая показывает, умеет ли он свести новое условие к уже решённому.

Хороший ответ занимает несколько минут и выглядит так:

1. Предложить жадность и сразу опровергнуть её на [2, 3, 2].
2. Сформулировать выбор для последнего дома и записать формулу.
3. Коротко сказать про экспоненту у рекурсии и $O(n)$ у мемоизации.
4. Написать версию с двумя переменными и объяснить, почему хватает двух значений.
5. Прогнать код на примере и на крайних случаях: один дом, два дома.

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

Что запомнить

• Смотрите на последний элемент: либо он в ответе, либо нет. Эта формулировка даёт рекурсию, остальное механика.
• Путь всегда один: рекурсия, мемоизация, таблица, сжатие памяти. Сжимать можно, когда формула смотрит на фиксированное число предыдущих значений.
• Таблица нужна, если просят показать сам выбор, а не только сумму.
• Круг разрывается перебором: что делаем с одним элементом. Здесь это два запуска линейного решения.