Тренировочная сессия 17: перекладка корня дерева СЕССИЯ

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

Финальная тренировка «Плато» — приём, который уже стоит на границе с «Пиком»: перекладка корня (её ещё называют «динамикой на дереве в две стороны»).

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

Обе задачи на этот приём, но переход в них разный. В первой перенос корня в соседа меняет ответ на $\pm 1$ — самая простая из возможных форм. Во второй приходится честно вычитать вклад поддерева и добавлять вклад верхней части, и там же выясняется, почему для максимума этот фокус не проходит без дополнительной подготовки.

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

---

Теория

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

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

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

Идея. Подвесим дерево за вершину 1. Для каждой вершины v посчитаем две величины:

• down[v] — ответ, если разрешено использовать только поддерево v (обычный обход снизу вверх);
• up[v] — ответ для «всего остального»: части дерева, которая остаётся, если отрезать поддерево v, но считая со стороны предка.

Тогда финальный ответ для v собирается из этих двух частей.

Первый проход — снизу вверх: down[v] считается из down дочерних вершин.

Второй проход — сверху вниз: up[u] для вершины u, дочерней по отношению к v, собирается из трёх слагаемых — собственного значения v, вклада остальных дочерних вершин v (то есть суммы по всем дочерним вершинам минус вклад самого u) и вклада верхней части up[v].

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

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

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

Обход, задающий order, при больших $n$ обязан быть нерекурсивным. И обратите внимание на строку down[v] += down[to] + sz[to]: вклад дочерней вершины входит суммой, поэтому его и можно вычесть при перекладке. Через максимум так не получится — там нужны два наибольших значения.

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

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

• ответ собирается суммой — перекладка пишется в две строки;
• ответ собирается максимумом или минимумом — придётся хранить два лучших значения на вершину (лучшее и второе лучшее), чтобы уметь исключать вклад одного потомка;
• ответ собирается чем-то принципиально необратимым — перекладка не подходит, нужен другой приём.

Три типовые формы перехода. Полезно узнавать их в лицо, потому что формула каждый раз одна и та же:

1. Вычесть и добавить. Ответ вершины — сумма вкладов потомков; при переносе корня из v в u мы вычитаем из ответа v вклад поддерева u и добавляем получившееся как «верхний» вклад для u. Так устроена сегодняшняя вторая задача.
2. Сдвиг на константу. Ответ меняется на одно и то же по форме число независимо от структуры поддерева — например, «плюс единица за одно ребро». Так устроена первая задача: перенос столицы меняет цену ровно на $\pm 1$.
3. Сдвиг, зависящий от размера поддерева. Классика — сумма расстояний до всех вершин: при переносе корня в соседа все вершины поддерева стали ближе на единицу, а все остальные — дальше на единицу, то есть ответ меняется на $n - 2 \cdot \text{размер поддерева}$.

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

Где приём перестаёт работать. Перекладка отвечает на вопрос «ответ для каждой вершины как для корня». Если ответ нужен только для одной вершины, второй проход не нужен вовсе. Если же ответ зависит не от корня, а от пары вершин (расстояния между всеми парами, путь между двумя запросами), перекладка не поможет — там нужны двоичные подъёмы и наименьший общий предок, к которым мы придём на «Пике».

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

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

Задача 1 — CF ~1700 — «Choosing Capital for Treeland»

Что дано

Дано n городов и n − 1 односторонняя дорога между ними. Если забыть про направления, из любого города можно добраться до любого другого — то есть перед нами дерево, у каждого ребра которого проставлена стрелка.

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

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

• Строка 1: целое n ($2 \le n \le 2 \cdot 10^5$).
• Каждая из следующих n − 1 строк: два числа s и t — дорога направлена из s в t.

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

• Строка 1: минимальное число разворотов.
• Строка 2: номера всех подходящих столиц в возрастающем порядке.

Пример 1:

Ввод:
3
2 1
2 3

Вывод:
0
2

Пример 2:

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

Вывод:
2
1 2 3

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

Возьмём второй пример: три дороги ведут в город 4 — из 1, из 2 и из 3. Это «звезда», центр которой всё собирает.

Посчитаем стоимость для каждой возможной столицы вручную.

• Столица 1. Из 1 в 4 доехать можно, дорога так и направлена. А вот до 2 и до 3 не добраться: дороги идут в обратную сторону. Разворачиваем две — цена 2.
• Столица 2. Симметрично: разворачиваем дороги к 1 и к 3 — цена 2.
• Столица 3. Так же, цена 2.
• Столица 4. Все три дороги ведут в четвёрку, значит выехать из неё нельзя ни по одной. Разворачиваем все три — цена 3.

Минимум равен 2 и достигается на городах 1, 2, 3 — совпадает с эталоном.

Теперь главное наблюдение. Считать так для каждого города по отдельности — это n обходов, то есть $O(n^2)$, а при $2 \cdot 10^5$ городах не проходит. Но посмотрите на цены: 2, 2, 2, 3. Они почти одинаковые, и переход от одного города к соседнему меняет цену ровно на единицу.

Это не совпадение. Пусть мы знаем ответ для города v и хотим ответ для соседнего u. Все дороги, кроме одной, ведут себя одинаково: путь от новой столицы к любому городу отличается от старого пути только началом. Меняется ровно одно ребро — то, что соединяет v и u:

• если дорога направлена из v в u, то при столице v её разворачивать не надо, а при столице u — надо: цена вырастает на 1;
• если дорога направлена из u в v, то наоборот: раньше мы за неё платили, теперь не платим — цена падает на 1.

Значит достаточно честно посчитать ответ для одного города, а все остальные получить, «перекладывая корень» по рёбрам.

Идея решения

Алгоритм в одну фразу: посчитать одним обходом стоимость для столицы в вершине 1, а затем вторым обходом получить стоимость для каждой вершины по формуле «стоимость предка ± 1 в зависимости от направления соединяющей дороги».

Как удобно закодировать направления. Каждую дорогу кладём в список смежности дважды, с весом: в направлении стрелки — вес 0 (ехать можно, платить не надо), против — вес 1 (придётся разворачивать). После этого дерево становится обычным неориентированным, а вся информация о стрелках живёт в весах.

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

Второй проход. Идём по вершинам в том же порядке, что и в первом обходе (важно: предок обязан быть обработан раньше потомка), и для каждого ребра v → u считаем

$\text{cost}[u] = \text{cost}[v] + \begin{cases} +1, & \text{ребро направлено из } v \text{ в } u \\ -1, & \text{ребро направлено из } u \text{ в } v \end{cases}$

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

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

Псевдокод

прочитать n и рёбра
для каждого ребра (s, t):
    добавить в список смежности s ребро (t, 0)      // по стрелке — бесплатно
    добавить в список смежности t ребро (s, 1)      // против стрелки — цена 1

// первый обход: цена для столицы в вершине 1
base = 0
обойти дерево из вершины 1, запоминая порядок обхода и предков:
    при переходе из v в u по ребру с весом w: base += w

cost[1] = base

// второй обход: перекладываем корень вдоль рёбер
для v в порядке обхода:
    для каждого ребра (u, w) из v, где u не предок v:
        cost[u] = cost[v] + (1 если w == 0 иначе -1)

best = минимум cost по всем вершинам
вывести best
вывести все вершины с cost, равным best, по возрастанию

Код решения

Комментарии по реализации. Оба обхода нерекурсивные: дерево из $2 \cdot 10^5$ вершин легко оказывается цепочкой, и рекурсия на такой глубине падает. Порядок order — это порядок посещения при обходе в глубину, и он гарантирует, что предок обработан раньше потомка; именно поэтому второй проход можно писать простым циклом по массиву, без нового обхода. Условие u != 1 во втором проходе защищает уже посчитанное значение корня от перезаписи, а u != parent[v] не даёт уйти вверх. Стоимость помещается в 32-битный тип (она не больше n − 1), но накопитель взят 64-битным — привычка, которая ничего не стоит и иногда спасает.

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

Пример 1: три города, дороги 2 → 1 и 2 → 3.

Списки смежности с весами: у вершины 1 — (2, 1); у вершины 2 — (1, 0), (3, 0); у вершины 3 — (2, 1).

Первый обход от вершины 1: спускаемся в 2 по ребру веса 1 (дорога направлена из 2 в 1, то есть против нашего движения), затем из 2 в 3 по ребру веса 0. Итого cost[1] = 1.

Минимум равен 0 и достигается только на вершине 2.

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

Пример 2: «звезда» из разбора выше — три дороги ведут в вершину 4.

Первый обход от вершины 1: ребро 1 → 4 имеет вес 0 (дорога направлена из 1 в 4), дальше из 4 в 2 и из 4 в 3 — веса 1 каждое (дороги направлены в четвёрку). Итого cost[1] = 2.

Минимум равен 2, достигается на вершинах 1, 2, 3.

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

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

• n = 2. Одна дорога; ответ 0, и подходит ровно тот город, из которого она выходит.
• Все дороги направлены от вершины 1. Тогда cost[1] = 0, и вершина 1 — единственный ответ; проверка, что первый проход считает базу правильно.
• Все дороги направлены в одну вершину (как во втором примере). Худшая столица — как раз этот «сток».
• Цепочка из $2 \cdot 10^5$ вершин. Проверка на нерекурсивность обхода.
• Несколько оптимальных столиц. Их надо вывести все и в возрастающем порядке — простой проход по вершинам от 1 до n даёт нужный порядок автоматически.
• Дерево-звезда с центром-источником. Все дороги выходят из центра: ответ 0 для центра и n − 2 для любого листа.

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

1. Запускать обход из каждой вершины. Логика верна, но это $O(n^2)$: на $2 \cdot 10^5$ вершинах решение не уложится.
2. Пересчитывать веса при перекладке в другую сторону. Знак определяется направлением исходной дороги, а не тем, куда мы идём: из v в u по «попутной» дороге цена растёт, по «встречной» — падает. Перепутать знаки — самая частая ошибка в этой задаче.
3. Обрабатывать вершины во втором проходе в произвольном порядке. Формула опирается на уже посчитанное значение предка; если предок ещё не готов, значения поедут по всему поддереву.
4. Строить дерево заново для каждой столицы. Перекладка корня не перестраивает структуру — она только пересчитывает число.
5. Рекурсивный обход. Падение на цепочке, хотя алгоритм правильный.
6. Вывести только один оптимальный город. Требуются все, и в возрастающем порядке.

Сложность

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

---

Задача 2 — CF ~1800 — «Maximum White Subtree»

Что дано

Дано дерево из n вершин. Каждая вершина покрашена: a[i] = 1 — белая, a[i] = 0 — чёрная.

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

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

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

Формат вывода: n чисел — ответы для вершин от 1 до n.

Пример 1:

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

Вывод:
2 2 2 2 2 1 1 0 2

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

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

Для примера значения получаются такими:

Дерево, подвешенное за вершину 1: дочерние вершины корня — 2 и 3; дочерняя вершина 2 — вершина 6, у неё дочерняя вершина 8; дочерние вершины 3 — вершины 4 и 5; дочерняя вершина 4 — вершина 7, дочерняя вершина 5 — вершина 9.

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

Ответ для вершины 1 — это down[1] = 2, и он совпадает с первым числом ожидаемого вывода.

Но для остальных вершин down — только половина ответа. Возьмём вершину 5: её down равен 0 (сама минус один, плюс белая девятка). Однако подграф, содержащий вершину 5, может уходить вверх — включить вершину 3 (+1) и её дочерней вершины 4 (+1). Это даёт 0 + 2 = 2 — и правильный ответ для вершины 5 равен 2, а не 0.

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

Идея решения

Алгоритм в одну фразу: посчитать down[v] обходом снизу вверх (значение вершины плюс положительные вклады дочерних вершин), затем обходом сверху вниз посчитать up[v] — лучший вклад верхней части, — и выдать down[v] + max(0, up[v]).

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

Первая часть — это down[v]. Формула down[v] = знач[v] + Σ max(0, down[u]) по дочерним вершинам u корректна, потому что каждое поддерево дочерней вершины присоединяется или нет независимо от остальных: связность сохраняется в любом случае, а вклад мы берём, только если он положителен.

Вторая часть — up[v]. Для вершины u, дочерней по отношению к v, она собирается так:

$up[u] = \text{знач}[v] + \Big(\sum_{w \text{ — дочерние вершины } v} \max(0, down[w]) - \max(0, down[u])\Big) + \max(0, up[v])$

То есть: сама вершина v, плюс вклады всех её дочерних вершин кроме u, плюс лучший вклад части, лежащей выше v. Вычитание здесь законно, потому что вклад суммируется, а не берётся максимумом.

Для корня верхней части нет: полагаем up[1] = 0, и тогда max(0, up[1]) = 0 ничего не добавляет.

Финальный ответ: down[v] + max(0, up[v]) — верхняя часть присоединяется только если её вклад положителен.

Псевдокод

прочитать n, цвета, рёбра; значение = +1 для белой, -1 для чёрной

обойти дерево из вершины 1 нерекурсивно, запомнив порядок и предков

// проход 1: снизу вверх
down[v] = значение[v] для всех v
для v в обратном порядке обхода:
    p = предок[v]
    если p != 0 и down[v] > 0:
        down[p] += down[v]

// проход 2: сверху вниз
up[1] = 0
для v в прямом порядке обхода:
    база = значение[v] + max(0, up[v])
    сумма_дочерних = down[v] - значение[v]        // это Σ max(0, down[дочерней])
    для каждого дочерней вершины u вершины v:
        up[u] = база + сумма_дочерних - max(0, down[u])

для v от 1 до n:
    вывести down[v] + max(0, up[v])

Код решения

Комментарии по реализации. Сумма положительных вкладов дочерних вершин не хранится отдельным массивом: она равна down[v] − знач[v] по построению первого прохода. Это экономит память и убирает шанс рассинхронизации двух величин. Оба прохода идут по одному и тому же массиву order: обратный порядок гарантирует, что дочерние вершины обработаны раньше предков, прямой — наоборот. Вывод собирается в одну строку: $2 \cdot 10^5$ отдельных печатей заметно медленнее.

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

Пример 1: значения и down посчитаны в разборе выше: down = [_, 2, 1, 2, 1, 0, −1, −1, −1, 1].

Считаем up сверху вниз:

Складываем down[v] + max(0, up[v]):

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

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

• n = 1. Дерево из одной вершины; ответ — её собственное значение (+1 или −1). Циклы по рёбрам не выполняются ни разу.
• Все вершины чёрные. Любой подграф только ухудшает сумму, поэтому ответ для каждой вершины равен −1 — подграф из неё одной.
• Все вершины белые. Ответ для каждой равен n: берём всё дерево.
• Чёрный корень, белые листья. Проверка того, что «верхняя» часть подключается только при положительном вкладе.
• Дерево-цепочка на $2 \cdot 10^5$ вершин. Тест на нерекурсивность обоих проходов.
• Вклад ровно ноль. Присоединять такое поддерево или нет — безразлично; max(0, ...) разрешает обе трактовки, ответ не меняется.

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

1. Запускать обход из каждой вершины. $O(n^2)$: на $2 \cdot 10^5$ вершин безнадёжно, хотя на примерах работает.
2. Забыть про max(0, ...). Присоединение поддерева с отрицательным вкладом обязательно делать необязательным; без обрезки ответ занижается.
3. Считать «сумму по остальным дочерним вершинам» отдельным циклом. Даёт $O(\text{степень}^2)$ и на «звезде» превращается в квадрат; вычитание одного слагаемого из общей суммы — правильный способ.
4. Вычитать вклад дочерней вершины там, где он входит максимумом. В этой задаче вклад суммируется, поэтому вычитание корректно; в задачах с максимумом нужны два наибольших значения.
5. Начать второй проход не с корня. Порядок обязан быть сверху вниз: up[u] использует уже посчитанное up[v] предка.
6. Рекурсивные проходы. Цепочка из $2 \cdot 10^5$ вершин кладёт стек.
7. Печатать ответы по одному. На больших n вывод становится узким местом.

Сложность

Время: $O(n)$ — два прохода по дереву. Память: $O(n)$.

---

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

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

• Codeforces 1092F «Tree with Maximum Cost» (~CF 1900) — у вершин есть веса, и стоимость дерева при корне v равна сумме «вес вершины, умноженный на расстояние до неё»; нужно выбрать корень с максимальной стоимостью. Каноническая перекладка: переход выражается через сумму весов поддерева, и формула получается ровно третьего типа из теории.
• Codeforces 337D «Book of Evil» (~CF 2000) — известны вершины, где обнаружено воздействие, и радиус его действия; нужно посчитать, сколько вершин могут быть источником. Здесь два прохода считают не сумму, а максимум расстояния до «заражённых» — хороший повод потренировать вариант с двумя лучшими значениями на вершину.
• Codeforces 1060E «Sergey and Subway» (~CF 2000) — нужно посчитать сумму расстояний между всеми парами вершин в дереве, где к каждой паре вершин на расстоянии два добавили ребро. Подсказка: сведите ответ к сумме обычных расстояний и чётности глубин, а сумму расстояний по всем парам считайте вкладом каждого ребра.
• Codeforces 1187E «Tree Painting» (~CF 2100) — вершины закрашиваются по одной, начиная с выбранной, и каждое закрашивание приносит размер связной белой компоненты, в которой оно происходит; нужно выбрать стартовую вершину с максимальным итогом. Самая сложная в наборе: сначала придётся доказать, что итог не зависит от порядка закрашивания, и только потом перекладывать корень.

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

---

Что дальше

«Плато» закончено: девять сессий, от жадности по разрядам до перекладки корня. Если оглянуться, видно общую линию всей фазы — почти каждая задача решалась не новым алгоритмом, а правильно выбранным состоянием или правильно доказанным свойством оптимума. Коридор ~CF 1400–1800 — это именно тот уровень, где перестаёт хватать знания приёмов и начинает требоваться умение их обосновывать.

Дальше — «Пик»: коридор ~CF 1900–2400, где приёмы уже не объясняются с нуля, а собираются в связки. Первая тренировка новой фазы — кратчайшие пути алгоритмом Дейкстры. До неё стоит закрыть хвосты: если какая-то из девяти сессий «Плато» осталась непрорешанной до конца, сейчас самое время вернуться — на «Пике» эти приёмы будут использоваться как готовые кирпичи, без повторного объяснения.