Задачник Базовый курс EN

Сколько раскрасок

дерево поиска дерево динамика подсчёт вариантов олимпиады

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

Раскраска — это выбор цвета, красного или чёрного, для каждого узла. Она правильная, если корень чёрный, у красного узла нет красных детей и на всех путях от корня вниз до пустых ссылок одинаковое число чёрных узлов. Ключи в узлах не важны, только форма. Третий, четвёртый, пятый и девятый тесты — формы настоящих деревьев: АВЛ-дерева из 19 073 землетрясений каталога USGS, вставленных по порядку времени, красно-чёрного и обычного дерева поиска из первых 20 000 разных слов «Войны и мира», АВЛ-дерева из адресов страниц legost.in. Номера узлов перемешаны.

Вход

В первой строке n и r (1 ≤ n ≤ 10⁵, 1 ≤ r ≤ n) — число узлов и номер корня. Дальше n строк: в i-й — номера левого и правого ребёнка узла i, 0 — если ребёнка нет. Гарантируется, что это двоичное дерево с корнем r, в котором есть все n узлов.

Выход

Число правильных раскрасок по модулю 10⁹ + 7.

Примеры

Пример 1

Вход

Выход

Пример 2

Вход

Выход

Ограничения

Время · Python: 4000 мс
Время · Go: 1000 мс
Память: 256 МБ
Проверка: 11 тестов

Разбор

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

1Подсказка 1

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

2Подсказка 2

Пусть у поддерева узла v чёрная высота b — столько чёрных узлов, считая v, на любом пути от v вниз. Если v чёрный, то у обоих детей, любого цвета, чёрная высота b − 1. Если красный — оба ребёнка чёрные, с чёрной высотой b. Пустая ссылка — чёрная, с чёрной высотой 0.

3Подсказка 3

Перебирать все b от 0 до n для каждого узла — слишком долго. Какие значения b вообще возможны, если самый короткий путь от v вниз до пустой ссылки содержит s узлов, а самый длинный — ℓ?

4Идея

Динамика по дереву снизу вверх. Для каждого узла v и чёрной высоты b храним два числа: B(v, b) — раскраски поддерева с чёрным v, R(v, b) — с красным. Тогда B(v, b) = (B + R)(лев., b − 1) · (B + R)(прав., b − 1), R(v, b) = B(лев., b) · B(прав., b), у пустой ссылки B = 1 при b = 0. Ответ — сумма B(корень, b) по всем b. Возможные b лежат между ⌊ℓ/2⌋ и s, а у сбалансированного дерева s не больше log₂(n + 1).

5Разбор

Почему формулы верны. Правила красно-чёрного дерева проверяются по поддеревьям: раскраска поддерева v правильна (не считая условия про корень), если правильны раскраски поддеревьев детей, у красного v нет красных детей и чёрные высоты детей равны. Чёрная высота v — это чёрная высота детей плюс один, если v чёрный. Раскраски левого и правого поддеревьев выбираются независимо, поэтому числа перемножаются, а разные цвета ребёнка — складываются. Условие «корень чёрный» учитывается в самом конце: суммируем только B.

Какие b возможны. На самом коротком пути от v вниз s узлов, и все чёрные узлы пути среди них, так что b ≤ s. На самом длинном пути ℓ узлов, красные не стоят подряд, и за каждым красным идёт чёрный (пустая ссылка в конце тоже чёрная), поэтому красных не больше ⌈ℓ/2⌉, а чёрных среди узлов пути не меньше ⌊ℓ/2⌋. Для каждого узла храним массив только для b от ⌊ℓ/2⌋ до s; если этот отрезок пуст, раскрасок нет. Когда ℓ больше 2s + 1 хоть у одного узла, ответ — ноль: так случается с обычным деревом поиска из слов романа: у его корня самый короткий путь вниз — 5 узлов, а самый длинный — 33.

Порядок обхода. Детей надо посчитать раньше родителя. Рекурсия на пути из 10⁵ узлов (десятый тест) переполнит стек Python, поэтому порядок строится явным стеком: обход сверху вниз, потом проход по нему в обратную сторону.

Что получается на настоящих формах. АВЛ-деревья раскрашиваются множеством способов — по модулю 10⁹ + 7 ответ выглядит случайным числом. Красно-чёрное дерево из слов романа, разумеется, раскрашивается: одну раскраску ему дала сама вставка. А самое худое АВЛ-дерево высоты 23, дерево Фибоначчи из 75 024 узлов, раскрашивается ровно одним способом: у каждого его узла самый длинный путь вниз почти вдвое длиннее самого короткого, и чёрную высоту выбирать не из чего.

5+Решение

6Сложность

Для узла хранится не больше s − ⌊ℓ/2⌋ + 1 значений, и каждое считается за O(1). В дереве из n узлов s у корня не больше log₂(n + 1) — иначе в верхних s этажах не хватило бы узлов, — и у любого узла тоже. Всего O(n log n) времени и памяти, а на деле много меньше: у большинства узлов поддеревья маленькие.

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

Красный под красным. Если у красного узла разрешить детей любого цвета, ответ получается больше правильного — уже на третьем тесте.
Красный корень. По условию корень чёрный; если сложить B и R у корня, ответ тоже завысится.
Рекурсия. Путь из 10⁵ узлов рекурсивный обход в Python не переживёт.
Все b подряд. Массив длины n на каждый узел — это 10¹⁰ операций; нужны только b между ⌊ℓ/2⌋ и s.

Алгоритмы задачи

DS Структура данных Красно-чёрное дерево Дерево поиска с узлами двух цветов: высота до 2·log₂(n + 1), не больше трёх поворотов на операцию. O(log n) на поиск, вставку и удаление 2 задачи DP Приём Динамика на деревьях Ответ поддерева — из ответов детей: проход от листьев к корню, а для всех корней сразу — перерасчёт с корня. O(n) 2 задачи