Тренировочная сессия 11: обход в глубину с состоянием вдоль пути СЕССИЯ

Что тренируем сегодня

Две предыдущие сессии были про массивы: отсортировать, пройти указателями, сопоставить. Сегодня структура меняется — вместо линейного массива дерево, и вместе с ней меняется способ думать. В массиве «соседний элемент» очевиден, в дереве соседей у вершины может быть сколько угодно, и вся работа сводится к тому, чтобы обойти всё ровно один раз и по дороге что-нибудь посчитать.

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

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

Коридор ~CF 1500–1600. Алгоритмически задачи несложные, но обе легко провалить на технике: неверном определении листа и обходе, который падает по глубине рекурсии.

---

Теория

Обход дерева в глубину с состоянием, накопленным вдоль пути. Спускаясь от корня, несём с собой одну величину, которая описывает весь пройденный путь; в каждой вершине она пересчитывается за один шаг из значения предка.

Когда применять:

• дано дерево, подвешенное за корень, и вопрос про пути от корня до вершин или листьев;
• состояние пути пересчитывается по одному шагу: зная значение в предке и данные текущей вершины, мы сразу знаем значение в ней (длина текущей серии, сумма, максимум, чётность);
• размер дерева до $10^5$–$10^6$, то есть за один обход надо ответить сразу для всех вершин.

Как решать:

1. Построить список смежности. Дерево обычно задано просто списком рёбер — направления в нём нет, корень выбираем сами (чаще всего вершина 1).
2. Запустить обход в глубину из корня, передавая вниз накопленное состояние.
3. Предка запоминаем, чтобы не уйти назад по тому же ребру: в дереве этого достаточно, отдельный массив «посещено» не нужен.
4. Если ограничение уже нарушено — не спускаться дальше. Это не оптимизация, а часть логики: у пути к любому потомку нарушение никуда не денется, префикс пути общий.

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

Про рекурсию: дерево из $10^5$ вершин может оказаться «бамбуком» — цепочкой, где глубина равна количеству вершин. Рекурсивный обход в такой ситуации падает по глубине стека (в Python — на первой же тысяче вложенных вызовов). Надёжный вариант — обход с явным стеком: кладём в стек кортеж «вершина, предок, состояние» и разбираем стек в цикле.

Каркас приёма (bad и step подставляются под конкретную задачу):

Какие бывают состояния. Приём один, но величина, которую тянут вниз, каждый раз своя. Полезно держать в голове весь набор:

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

Отдельно про максимум по предкам. Если веса рёбер неотрицательны, «самый длинный путь от какого-нибудь предка» — это просто расстояние от корня. Но стоит появиться отрицательным весам, и наивная формула ломается: далёкий предок может дать меньшую сумму, чем близкий. Правильный пересчёт — «либо путь начинается прямо в предке, либо продолжает лучший путь, который в предке заканчивался»:

best[v] = вес_ребра(предок, v) + max(0, best[предок])

Ноль под максимумом — это и есть вариант «начать путь заново в предке». Та же самая формула, что в задаче о максимальной сумме подотрезка массива, только вместо массива путь дерева.

Почему обрезать ветку целиком — законно. Отсечение здесь не оптимизация, а часть логики. Путь к любому потомку вершины v содержит путь к v целиком, как префикс. Если ограничение нарушено уже на этом префиксе, у потомков оно нарушено тем более — заходить внутрь незачем и вредно: можно случайно засчитать лежащий там лист. Ровно поэтому проверка стоит до разворачивания соседей, а не после.

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

С чем путают. Оба обхода — в глубину, различаются только моментом вычисления. Простое правило: значение считается до спуска к потомкам — состояние идёт вдоль пути сверху; значение считается после возврата от потомков — это уже динамика по поддеревьям. Если в задаче нужны обе половины, обычно пишут один обход и два места вычисления, а не два обхода.

Сложность приёма: $O(n)$ — каждое ребро просматривается дважды.

---

Задача 1 — CF ~1500 — «Kefa and Park»

Что дано

Дано дерево из n вершин, подвешенное за вершину 1. В некоторых вершинах живут коты: a[i] = 1, если в вершине i есть кот, и a[i] = 0 иначе. Пункты назначения расположены в листьях дерева.

Путь считается допустимым, если на нём не встречается более m подряд идущих вершин с котами. Нужно посчитать, сколько листьев достижимо по допустимому пути из вершины 1.

Формат ввода:

• Строка 1: два целых числа n и m ($2 \le n \le 10^5$, $1 \le m \le n$).
• Строка 2: n чисел a[1] … a[n], каждое 0 или 1.
• Каждая из следующих n − 1 строк: два числа — концы ребра дерева.

Формат вывода: одно число — количество достижимых листьев.

Пример 1:

Ввод:
4 1
1 1 0 0
1 2
1 3
1 4

Вывод:
2

Пример 2:

Ввод:
7 1
1 0 1 1 0 0 0
1 2
1 3
2 4
2 5
3 6
3 7

Вывод:
2

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

Возьмём пример 1. Дерево — «звезда»: корень 1 и три листа 2, 3, 4. Коты стоят в вершинах 1 и 2, ограничение m = 1, то есть двух котов подряд уже достаточно, чтобы путь стал недопустимым.

Ответ 2. Обратите внимание: считается не общее количество котов на пути, а длина непрерывной серии — вершина без кота обнуляет счётчик.

Теперь пример 2. Дерево двухуровневое: у корня 1 два дочерней вершины (2 и 3), у каждого из них по два дочерней вершины. Коты — в вершинах 1, 3, 4.

Ответ 2. Видно и главное следствие для алгоритма: как только на пути к вершине 3 серия достигла двух, обе ветки под ней отпадают целиком — заходить в них незачем. Путь к любому потомку проходит через уже испорченный префикс.

Идея решения

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

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

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

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

Псевдокод

прочитать n, m, массив a, рёбра
построить список смежности g

ответ = 0
стек = [ (вершина 1, предок 0, серия a[1]) ]

пока стек не пуст:
    (v, p, серия) = снять со стека

    если серия > m:
        продолжить          // ветка обрезана: ниже все листья недостижимы

    лист = да
    для каждого to из g[v]:
        если to == p: пропустить
        лист = нет
        если a[to] == 1: новая_серия = серия + 1
        иначе:           новая_серия = 0
        положить в стек (to, v, новая_серия)

    если лист:
        ответ = ответ + 1

вывести ответ

Код решения

Комментарии по реализации. Обход написан через явный стек в обоих языках: в Python рекурсия на цепочке из $10^5$ вершин упирается в лимит глубины, в C++ — рискует переполнить стек потока. Массив «посещено» не заводится: в дереве нет циклов, поэтому запрета возвращаться к предку достаточно, чтобы каждая вершина попала в стек ровно один раз. Значение run для корня инициализируется как a[1], а не нулём, — кот в корне тоже начинает серию.

Проверка на примерах

Пример 1: n = 4, m = 1, a = [_, 1, 1, 0, 0].

Наш вывод: 2 — совпадает с ожидаемым.

Пример 2: n = 7, m = 1, a = [_, 1, 0, 1, 1, 0, 0, 0].

Наш вывод: 2 — совпадает с ожидаемым. Обратите внимание, что вершины 6 и 7 в обход не попадают вообще: их предок 3 отсечён.

Крайние случаи

• n = 2. Корень и один лист. Корень листом не считается, ответ 0 или 1 в зависимости от котов.
• Кот в корне и m = 1. Серия стартует с единицы; любой дочерняя вершина с котом уже недостижим.
• Котов нет вообще. Все листья достижимы; ответ равен количеству листьев.
• Коты во всех вершинах. Достижимы только листья на глубине не больше m рёбер от корня.
• m ≥ n. Ограничение недостижимо, ответ — количество всех листьев.
• Дерево-цепочка длины $10^5$. Единственный лист — конец цепочки; тест на переполнение стека рекурсии.
• Корень степени 1. Классическая ловушка: корень выглядит «как лист» по степени, но листом не является.

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

1. Считать корень листом по степени. Проверка «степень равна 1» без исключения для корня даёт лишнюю единицу в ответе на дереве-цепочке.
2. Копить общее число котов вместо длины непрерывной серии. Ограничение именно на подряд идущие вершины; вершина без кота обнуляет счётчик.
3. Не обрезать ветку. Если после превышения продолжать спуск и проверять «текущую серию» в листе, отсечённые листья могут ошибочно засчитаться: серия к моменту прихода в лист успевает обнулиться.
4. Рекурсивный обход в Python. На цепочке из $10^5$ вершин решение падает по глубине рекурсии, а не по времени.
5. Забыть про кота в корне. Инициализация серии нулём вместо a[1] даёт неверный ответ на всех тестах, где корень с котом.
6. Заводить массив «посещено» и проверять его вместо предка. Работает, но лишняя память и лишний источник ошибок; в дереве достаточно не возвращаться в предка.

Сложность

Время: $O(n)$ — каждое ребро просматривается не более двух раз. Память: $O(n)$ — список смежности и стек обхода.

---

Задача 2 — CF ~1600 — «Alyona and the Tree»

Что дано

Дано дерево с корнем в вершине 1. На каждой вершине v написано число a[v], на каждом ребре — вес (он может быть отрицательным). Вершина v называется грустной, если в её поддереве найдётся вершина u, для которой расстояние от v до u (сумма весов рёбер на пути) строго больше a[u].

Разрешается удалять листья — по одному, сколько угодно раз; после удаления листа его сосед может сам стать листом. Нужно найти минимальное число удалённых вершин, после которого в дереве не останется ни одной грустной вершины.

Формат ввода:

• Строка 1: целое n ($1 \le n \le 10^5$) — количество вершин.
• Строка 2: n чисел a[1] … a[n] ($1 \le a_i \le 10^9$).
• Каждая из следующих n − 1 строк: два числа p и c ($1 \le p \le n$, $-10^9 \le c \le 10^9$) — i-я по счёту строка задаёт ребро между вершиной i + 1 и вершиной p с весом c.

Формат вывода: одно число — минимальное количество удалённых вершин.

Пример 1:

Ввод:
9
88 22 83 14 95 91 98 53 11
3 24
7 -8
1 67
1 64
9 65
5 12
6 -80
3 8

Вывод:
5

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

Сначала распутаем ввод. Строки рёбер описывают вершины 2, 3, …, 9 по порядку: у вершины 2 предок 3 (вес 24), у вершины 3 предок 7 (вес −8), у вершины 4 предок 1 (вес 67), у 5 предок 1 (вес 64), у 6 предок 9 (вес 65), у 7 предок 5 (вес 12), у 8 предок 6 (вес −80), у 9 предок 3 (вес 8).

Собираем дерево от корня: у вершины 1 потомки 4 и 5; у 5 потомок 7; у 7 потомок 3; у 3 потомки 2 и 9; у 9 потомок 6; у 6 потомок 8.

Теперь переформулируем условие. Оно сформулировано «сверху вниз»: вершина v грустит из-за какого-то потомка u. Но удаляем-то мы вершины, а не пары, поэтому удобнее посмотреть с другого конца — снизу вверх: вершина u виновата, если найдётся хотя бы один её предок v, до которого расстояние больше a[u]. Тогда u придётся удалить. Значит для каждой вершины достаточно знать одно число: максимум расстояния от неё до какого-нибудь предка.

Посчитаем это число, спускаясь от корня. Обозначим его best[v].

• Вершина 4: единственный предок — корень, расстояние 67. a[4] = 14, и 67 > 14 — вершина 4 виновата.
• Вершина 5: расстояние до корня 64, a[5] = 95 — норма.
• Вершина 7: до вершины 5 расстояние 12, до корня 76. Максимум 76, a[7] = 98 — норма.
• Вершина 3: ребро в неё весит −8, поэтому до вершины 7 расстояние −8, а до корня 68. Максимум 68, a[3] = 83 — норма.
• Вершина 2: ребро 24, лучший путь через 3 даёт 68 + 24 = 92. a[2] = 22, 92 > 22 — виновата.
• Вершина 9: 68 + 8 = 76, a[9] = 11 — виновата.

Итого виноваты вершины 4, 2 и 9. Но удалить придётся не три вершины: если убрать вершину 9, её потомки 6 и 8 повисают в воздухе — их тоже надо удалить, и порядок «сначала листья» это допускает (сначала лист 8, потом 6, потом 9). Значит удаляется вся ветка: 9, 6, 8 — три вершины, плюс 4 и 2. Всего 5, как в эталоне.

Обратите внимание на вершину 3: ребро в неё отрицательное, и расстояние от её непосредственного предка равно −8 — то есть меньше нуля. Если бы мы считали best как «расстояние до корня» или как «максимум предыдущего плюс вес», получили бы неверные числа у неё и у её потомков.

Идея решения

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

Почему условие можно перевернуть. В исходной формулировке грустит предок, а удаляют потомков — неудобно. Но обе стороны говорят про одну и ту же пару «предок v — потомок u» с одним и тем же неравенством dist(v, u) > a[u]. Значит дерево свободно от грусти тогда и только тогда, когда в нём не осталось ни одной такой пары, а это условие естественно вешается на нижний конец: вершина u допустима, если ни один её предок не даёт слишком большого расстояния, то есть если максимум по предкам не превышает a[u].

Как считается максимум по предкам. Это тот самый случай из теории, где отрицательные веса ломают наивную формулу. Правильный пересчёт:

$\text{best}[v] = c(\text{предок}, v) + \max(0,\ \text{best}[\text{предок}])$

Разбор двух вариантов: путь до v либо начинается прямо в предке (тогда его длина — вес одного ребра), либо продолжает какой-то путь, заканчивавшийся в предке (тогда выгодно взять самый длинный из них). Если все пути, заканчивающиеся в предке, отрицательны, выгоднее начать заново — за это и отвечает ноль под максимумом. Для корня best равен нулю: предков у него нет, и формула для его потомков корректно даёт вес ребра.

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

Почему не нужно проверять, не станет ли кто-то грустным после удалений. Удаления только уменьшают дерево. Ни одна пара «предок — потомок» от этого не появляется, а значит новых нарушений возникнуть не может.

Псевдокод

прочитать n, массив a, рёбра (предок, вес) для вершин 2..n
построить списки смежности

удалено = 0
стек = [(1, 0, 0, ложь)]           // вершина, предок, best, «уже режем»

пока стек не пуст:
    v, предок, best, режем = снять со стека
    если не режем и best > a[v]:
        режем = истина            // отсюда и ниже всё поддерево уходит
    если режем:
        удалено += 1
    для каждого соседа u вершины v, кроме предка:
        положить в стек (u, v, вес(v, u) + max(0, best), режем)

вывести удалено

Код решения

Комментарии по реализации. Обход написан явным стеком, а не рекурсией: дерево из $10^5$ вершин запросто оказывается цепочкой, и рекурсивная версия падает по глубине — в Python на первой же тысяче вызовов, в C++ чуть позже, но так же неприятно. Веса до $10^9$ и путь до $10^5$ рёбер дают сумму до $10^{14}$, поэтому и best, и a — 64-битные; int здесь переполнится незаметно. Флаг cut избавляет от отдельного прохода «посчитать размер поддерева»: вершины срезанной ветки просто досчитываются тем же обходом.

Проверка на примерах

Пример 1: дерево из разбора выше.

Удалены: 4, 2, 9, 6, 8 — итого 5.

Наш вывод: 5 — совпадает с ожидаемым.

Полезно заметить, что вершины 6 и 8 сами по себе невиновны: у вершины 6 максимум по предкам был бы 65 + max(0, 76) = 141, что действительно больше a[6] = 91, а вот у вершины 8 ребро −80 сделало бы её вполне законной. Но она всё равно уходит — потому что висит под удалённой веткой, а листья снимаются только сверху донизу.

Крайние случаи

• n = 1. Рёбер нет, корень никем не может быть обижен, ответ 0. Заодно проверка, что чтение не ждёт лишних строк.
• Все веса отрицательные. Тогда best у каждой вершины равен весу единственного ребра к предку (максимум с нулём обнуляет накопленное), и, поскольку a[i] ≥ 1, удалять не придётся никого — ответ 0.
• Отрицательное ребро в середине пути. Главный случай, ради которого нужна формула с max(0, …): как вершина 3 в примере.
• Виновата вершина, соседняя с корнем. Отрезается сразу целое поддерево — проверка, что счётчик считает не только саму вершину.
• Цепочка из $10^5$ вершин с большими весами. Проверка и на глубину стека, и на 64-битную арифметику: сумма пути доходит до $10^{14}$.
• Виноватая вершина внутри уже срезанного поддерева. Она не должна считаться дважды — за это отвечает проверка if not cut перед сравнением.

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

1. Считать best как расстояние от корня. Работает только при неотрицательных весах. Одно отрицательное ребро — и ответ поедет.
2. Писать best[v] = max(best[предок] + c, 0). Похоже на правильную формулу, но ноль тут не на месте: обнулять надо накопленное до прибавления веса, а не после. Иначе путь длины −5 превратится в 0 и вершина ошибочно окажется законной.
3. Рекурсивный обход. На «бамбуке» из $10^5$ вершин решение падает, хотя логика верна.
4. Останавливать обход на виноватой вершине. Тогда её поддерево не посчитается, и ответ окажется меньше настоящего: удалить-то придётся всех.
5. Проверять виноватость внутри срезанного поддерева. Вершины там могут быть «законными», но это уже неважно — они всё равно удаляются. Лишняя проверка занижает ответ.
6. 32-битная арифметика. Сумма весов до $10^{14}$ не помещается в int, а переполнение здесь тихое.
7. Считать, что вершина i + 1 в строке ввода — это вершина i. Классическая ошибка на единицу в формате: строки описывают вершины начиная со второй.

Сложность

Время: $O(n)$ — каждое ребро просматривается дважды, обход один. Память: $O(n)$ — список смежности и стек, в худшем случае глубиной во всё дерево.

---

Самостоятельная тренировка

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

• Codeforces 1830A «Copil Copac Draws Trees» (~CF 1400) — рёбра дерева даны в некотором порядке и рисуются проходами: за один проход рисуются все рёбра, у которых один конец уже нарисован, причём в порядке их номеров; нужно узнать, за сколько проходов дерево будет нарисовано целиком. Подсказка: ответ для вершины считается вдоль пути от корня по одному правилу за шаг — ровно как длина серии в сегодняшней первой задаче.
• Codeforces 839C «Journey» (~CF 1500) — из корня дерева идут случайным образом, каждый раз равновероятно выбирая любую ещё не посещённую соседнюю вершину, и останавливаются, когда идти некуда; нужно найти математическое ожидание длины маршрута. Состояние вдоль пути здесь — накопленная вероятность попасть в текущую вершину; она умножается на шаге, а не складывается.
• Codeforces 1611E1 «Escape The Maze (easy version)» (~CF 1700) — в дереве расставлены преследователи; нужно понять, можно ли добраться из корня до какого-нибудь листа, ни разу не столкнувшись. Подсказка: спускаясь, полезно тянуть с собой расстояние до ближайшего преследователя над текущей вершиной и сравнивать его с собственной глубиной.
• Codeforces 1714G «Path Prefixes» (~CF 1700) — на рёбрах дерева написаны две величины сразу; для каждой вершины нужно найти самый длинный префикс пути от корня, у которого сумма по второй величине не превышает сумму по первой. Самая техничная в наборе: состояние вдоль пути здесь не число, а весь массив префиксных сумм, и по нему приходится искать бинарным поиском.

Прорешай их самостоятельно — если застрянешь, разбери с AI-партнёром на codepal.ru: он не даёт готовый ответ, а подводит к решению вопросами.

---

Что дальше

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

На следующей тренировке возвращаемся к массивам: скользящее окно по отсортированным данным — приём, у которого оба указателя идут в одну сторону, в отличие от встречных с сессии 10.