В этой статье рассматривается понятие *виртуального дерева*, доказываются его ключевые свойства, описывается эффективный алгоритм его построения, а затем метод применяется для решения конкретной задачи с использованием динамического программирования (ДП) на деревьях.↵
↵
## Введение↵
↵
Во многих задачах, связанных с деревьями, требуется работать с подмножеством вершин, сохраняя при этом структуру дерева, индуцированную их попарными наименьшими общими предками (LCA). Понятие *виртуального дерева* позволяет перейти от исходного дерева $T$ с $N$ вершинами к подструктуре, размер которой линейно зависит от размера выбранного множества $X$. Это существенно ускоряет алгоритмы, в частности, для подсчета ДП.↵
↵
## Определение виртуального дерева↵
↵
Пусть $T$ — заданное корневое дерево, а $X$ — некоторое подмножество его вершин. **Виртуальное дерево** $A(X)$ определяется следующим образом:↵
$$↵
A(X) = \{ \operatorname{lca}(x,y) \mid x,y \in X \},↵
$$↵
где $\operatorname{lca}(x,y)$ означает наименьшего общего предка вершин $x$ и $y$ в дереве $T$. В этом дереве между каждой парой вершин проводится ребро, если одна из них является предком другой в $T$.↵
↵
Такое определение гарантирует, что в структуру попадают все вершины, «важные» для анализа $X$ (то есть все LCA для пар вершин из $X$), а связность наследуется от исходного дерева.↵
↵
## Ключевые свойства и их доказательства↵
↵
### Свойство 1: Оценка по размеру↵
↵
**Утверждение.** ↵
Для любого множества $X$ вершин дерева $T$ справедливо неравенство:↵
$$↵
\vert A(X) \vert \le 2 \vert X \vert - 1.↵
$$↵
↵
**Доказательство.** ↵
Выберем порядок обхода в глубину (DFS) и рассмотрим вершины множества $X$, расположенные в порядке обхода: ↵
$$↵
x_1, x_2, \dots, x_m.↵
$$↵
Обозначим↵
$$↵
l = \operatorname{lca}(x_1, x_2).↵
$$↵
Далее необходимо доказать следующую лемму.↵
↵
> **Лемма 1.** Для любого целого числа $n \ge 3$, если $\operatorname{lca}(x_1,x_n) \neq l$, то↵
$$↵
\operatorname{lca}(x_1,x_n) = \operatorname{lca}(x_2,x_n).↵
$$↵
>↵
> **Доказательство (от противного).** ↵
> Предположим, что для некоторого $n \ge 3$ выполняется $\operatorname{lca}(x_1,x_n) \neq l$ и одновременно $\operatorname{lca}(x_1,x_n) \neq \operatorname{lca}(x_2,x_n)$. Пусть $k = \operatorname{lca}(x_1,x_n)$. ↵
> Если $k$ является предком $l$, то по определению получаем, что и $\operatorname{lca}(x_1,x_n)$, и $\operatorname{lca}(x_2,x_n)$ равны $k$, что противоречит предположению. Следовательно, $k$ должно быть потомком $l$. Но тогда в порядке DFS должен получиться порядок $x_1, x_n, x_2$, что противоречит выбранному порядку обхода. ↵
> Противоречие завершает доказательство леммы.↵
↵
Возвращаясь к доказательству свойства, заметим, что согласно лемме, каждый LCA, возникающий при последовательном рассмотрении вершин $X$, равен либо:↵
1. самой вершине $x_1$,↵
2. $\operatorname{lca}(x_1,x_2)$, либо↵
3. $\operatorname{lca}(x_2,x_n)$ для некоторого $n$.↵
↵
Таким образом, можно рекурсивно записать:↵
$$↵
A(\{x_1,\dots,x_m\}) = A(\{x_2,\dots,x_m\}) \cup \{ x_1, \operatorname{lca}(x_1,x_2) \}.↵
$$↵
Из этого следует, что↵
$$↵
\vert A(X) \vert \le \vert A(\{x_2,\dots,x_m\}) \vert + 2.↵
$$↵
Применяя индукцию по размеру множества $X$, получаем итоговую оценку:↵
$$↵
\vert A(X) \vert \le 2 \vert X \vert - 1.↵
$$↵
При этом, если дерево $T$ является совершенным двоичным деревом, а $X$ состоит из его листьев, неравенство достигается в точности.↵
↵
### Свойство 2: Представление через последовательные LCA↵
↵
**Утверждение.** ↵
Пусть вершины $X$ упорядочены по порядку обхода DFS: $x_1, x_2, \dots, x_m$. Тогда↵
$$↵
A(X) = X \cup \{ \operatorname{lca}(x_i, x_{i+1}) \mid 1 \le i < m \}.↵
$$↵
↵
**Доказательство.** ↵
Начнём с рекурсивного представления:↵
$$↵
A(\{x_1,\dots,x_m\}) = A(\{x_2,\dots,x_m\}) \cup \{ x_1, \operatorname{lca}(x_1,x_2) \}.↵
$$↵
Раскроем его рекурсивно:↵
$$↵
\begin{aligned}↵
A(\{x_1,\dots,x_m\}) &= A(\{x_2,\dots,x_m\}) \cup \{ x_1, \operatorname{lca}(x_1,x_2) \} \\↵
&= A(\{x_3,\dots,x_m\}) \cup \{ x_1, \operatorname{lca}(x_1,x_2), x_2, \operatorname{lca}(x_2,x_3) \} \\↵
&\quad \vdots \\↵
&= \{ x_1, \operatorname{lca}(x_1,x_2), x_2, \dots, x_{m-1}, \operatorname{lca}(x_{m-1},x_m), x_m \}.↵
\end{aligned}↵
$$↵
Таким образом, получаем требуемое представление.↵
↵
## Построение виртуального дерева↵
↵
Дано множество вершин $X$. Виртуальное дерево $A(X)$ можно построить за время↵
$$↵
O\left(|X| (\log |X| + \log N)\right),↵
$$↵
при этом необходимо выполнить предпосчет исходного дерева $T$ за $O(N \log N)$ для обеспечения быстрых запросов LCA.↵
↵
**Основные шаги построения:**↵
↵
1. **Предподсчет.** ↵
С помощью обхода в глубину (DFS) или методов, таких как двоичные подъемы или HLD, вычисляем времена входа в вершины и LCA для любых двух вершин.↵
↵
2. **Сортировка вершин.** ↵
Сортируем вершины множества $X$ в порядке их появления при обходе DFS (используя время входа).↵
↵
3. **Добавление LCA.** ↵
Для последовательных вершин $x_i$ и $x_{i+1}$ вычисляем $\operatorname{lca}(x_i,x_{i+1})$. Объединение $X$ и этих LCA даёт множество вершин $A(X)$ (согласно свойству 2).↵
↵
4. **Построение дерева.** ↵
Получив множество вершин $A(X)$, отсортированных по порядку DFS, можно пройти по последовательности, используя стек, для восстановления структуры дерева. При обходе поддерживается стек, верхний элемент которого является текущей вершиной, а если новая вершина не является потомком вершины на вершине стека, производится «выталкивание» элементов до нахождения соответствующего предка, после чего производится соединение.↵
↵
Такой алгоритм гарантирует, что виртуальное дерево содержит $O(|X|)$ вершин, что позволяет эффективно применять ДП.↵
↵
## Применение: подсчет поддеревьев в раскрашенном дереве↵
↵
### Условие [задачи](https://atcoder.jp/contests/abc340/tasks/abc340_g)↵
↵
Дано дерево $T$ с $N$ вершинами, пронумерованными от $1$ до $N$. Ребро $i$ соединяет вершины $u[i]$ и $v[i]$. Каждая вершина $i$ раскрашена в цвет $A[i]$. ↵
Необходимо найти (по модулю 998244353) количество (непустых) подмножеств $S$ вершин дерева $T$, удовлетворяющих условию:↵
- Индуцированный граф $G[S]$ является деревом.↵
- Все вершины $G[S]$ со степенью 1 имеют один и тот же цвет.↵
↵
### Основная идея решения↵
↵
Основная идея заключается в разбиении задачи по цветам. Для каждого цвета $c$ рассматривается множество вершин $X$, раскрашенных в $c$, и строится виртуальное дерево для $X$. Затем на полученном дереве выполняется ДП для подсчета допустимых поддеревьев, у которых все вершины со степенью 1 имеют цвет $c$. Итоговый ответ получается суммированием результатов по всем цветам.↵
↵
Благодаря построению виртуального дерева, хотя исходное дерево содержит $N$ вершин, каждое виртуальное дерево для конкретного цвета имеет размер $O(|X|)$, что позволяет выполнять DP за приемлемое время.↵
↵
Получаем, что суммарная сложность предпосчета и построения виртуальных деревьев не превышает $O(N \log N + \sum_{c \in C} |X_c| (\log |X_c| + \log N)) = O(N \log N)$, где $C$ это множество различных цветов вершин и $X_c$ это множество вершин с цветом $c$.↵
↵
### Реализация решения↵
↵
Ниже приведён полный код на C++ с комментариями, описывающими основные компоненты алгоритма:↵
↵
```cpp↵
#include <bits/stdc++.h>↵
using namespace std;↵
↵
const int MAXN = 2e5 + 10;↵
const int LOGN = 20;↵
const int MOD = 998244353;↵
↵
vector<int> g[MAXN]; // Список смежности заданного графа↵
vector<int> ver[MAXN]; // Для каждого цвета ver[c] хранит вершины с цветом c↵
vector<int> aux_g[MAXN]; // Список смежности для виртуального дерева↵
int col[MAXN]; // Цвет каждой вершины↵
int tmr, n; // Глобальный счетчик времени и количество вершин↵
int up[LOGN][MAXN], dep[MAXN], tin[MAXN]; // Для подсчета LCA и времени входа↵
↵
// DP массивы для динамического программирования на виртуальном дереве↵
int dp[MAXN][2], sum[MAXN];↵
↵
// Функция предподсчета: DFS для вычисления tin, массива up и dep↵
void dfs_precalc(int v, int p) {↵
tin[v] = ++tmr;↵
up[0][v] = p;↵
dep[v] = dep[p] + 1;↵
for (int i = 1; i < LOGN; ++i)↵
up[i][v] = up[i - 1][up[i - 1][v]];↵
for (auto to : g[v])↵
if (to != p)↵
dfs_precalc(to, v);↵
}↵
↵
// Функция для вычисления LCA с использованием двоичных подъемов↵
int getlca(int x, int y) {↵
if (dep[x] < dep[y]) swap(x, y);↵
for (int i = LOGN - 1; i >= 0; --i)↵
if (dep[up[i][x]] >= dep[y])↵
x = up[i][x];↵
if (x == y) return x;↵
for (int i = LOGN - 1; i >= 0; --i)↵
if (up[i][x] != up[i][y]) {↵
x = up[i][x];↵
y = up[i][y];↵
}↵
return up[0][x];↵
}↵
↵
// DFS по виртуальному дереву для выполнения DP.↵
// Параметр c — целевой цвет, для которого производится подсчет.↵
↵
void dfs_calc(int v, int p, int c, int &ans) {↵
dp[v][0] = dp[v][1] = 0;↵
sum[v] = 0;↵
for(auto to : aux_g[v]) {↵
if(to == p) continue;↵
dfs_calc(to, v, c, ans);↵
// Переходы DP: объединение текущего состояния с результатом из поддерева.↵
int nxt0 = (dp[v][0] + sum[to]) % MOD;↵
int nxt1 = ((dp[v][0] + dp[v][1]) * 1ll * sum[to] % MOD + dp[v][1]) % MOD;↵
dp[v][0] = nxt0;↵
dp[v][1] = nxt1;↵
}↵
sum[v] = (dp[v][0] + dp[v][1]) % MOD;↵
if(col[v] == c) {↵
// Если вершина имеет целевой цвет, она может участвовать в корректном поддереве.↵
sum[v] = (sum[v] + 1) % MOD;↵
ans = (ans + sum[v]) % MOD;↵
} else {↵
ans = (ans + dp[v][1]) % MOD;↵
}↵
}↵
↵
// Функция для построения виртуального дерева для цвета c и выполнения DP.↵
void calc_aux(int c, int &ans) {↵
auto p = ver[c];↵
if (p.empty()) return;↵
// Сортируем вершины по времени входа (tin) — порядку обхода DFS.↵
sort(p.begin(), p.end(), [&](const int a, const int b) { return tin[a] < tin[b]; });↵
vector<int> verstk = {1}; // Инициализируем стек с корнем дерева T (вершина 1).↵
aux_g[1].clear();↵
auto add = [&](int u, int v) {↵
aux_g[u].push_back(v);↵
aux_g[v].push_back(u);↵
};↵
// Обрабатываем каждую вершину из множества p, поддерживая стек для построения виртуального дерева.↵
for (auto u : p) {↵
if (u == 1) continue;↵
int lca = getlca(u, verstk.back());↵
if (lca != verstk.back()) {↵
while (verstk.size() >= 2 && tin[lca] < tin[verstk[verstk.size() - 2]]) {↵
add(verstk.back(), verstk[verstk.size() - 2]);↵
verstk.pop_back();↵
}↵
if (verstk.size() >= 2 && tin[lca] != tin[verstk[verstk.size() - 2]]) {↵
aux_g[lca].clear();↵
add(verstk.back(), lca);↵
verstk.back() = lca;↵
} else {↵
add(verstk.back(), lca);↵
verstk.pop_back();↵
}↵
}↵
aux_g[u].clear();↵
verstk.push_back(u);↵
}↵
while (verstk.size() > 1) {↵
add(verstk.back(), verstk[verstk.size() - 2]);↵
verstk.pop_back();↵
}↵
// Выполняем DP по виртуальному дереву, начиная с корня 1.↵
return dfs_calc(1, 0, c, ans);↵
}↵
↵
int main() {↵
ios_base::sync_with_stdio(0);↵
cin.tie(0);↵
cin >> n;↵
// Считываем цвета и группируем вершины по цвету.↵
for (int i = 1; i <= n; ++i) {↵
cin >> col[i];↵
ver[col[i]].push_back(i);↵
}↵
for (int i = 1; i < n; ++i) {↵
int u, v;↵
cin >> u >> v;↵
g[u].push_back(v);↵
g[v].push_back(u);↵
}↵
dfs_precalc(1, 0);↵
int ans = 0;↵
// Для каждого возможного цвета обрабатываем соответствующее виртуальное дерево.↵
for (int i = 1; i <= n; ++i)↵
calc_aux(i, ans);↵
cout << ans;↵
return 0;↵
}↵
```↵
↵
### Объяснение реализации↵
↵
- **Предподсчет.** ↵
Функция `dfs_precalc` выполняет обход дерева $T$ в глубину, вычисляя время входа (`tin`), глубину и заполняя таблицу двоичных подъемов `up` для быстрых запросов LCA.↵
↵
- **Вычисление LCA.** ↵
Функция `getlca` реализует двоичные подъемы для быстрого нахождения наименьшего общего предка двух вершин.↵
↵
- **Построение виртуального дерева.** ↵
Функция `calc_aux` для каждого цвета $c$ берёт множество вершин `ver[c]`, сортирует его по `tin` и с помощью стека строит виртуальное дерево. При этом для каждой пары последовательных вершин вычисляется LCA, что соответствует свойству 2.↵
↵
- **Динамическое программирование.** ↵
Функция `dfs_calc` выполняет обход виртуального дерева и объединяет результаты поддеревьев согласно переходам DP. Состояния DP рассчитаны таким образом, что учитывается вклад вершины, если она имеет целевой цвет, и производится подсчет только тех поддеревьев, где все листья имеют один цвет.↵
↵
- **Сбор результата.** ↵
Функция `main` считывает входные данные, выполняет предподсчет, и затем суммирует результаты для каждого цвета, выводя итоговый ответ по модулю 998244353.↵
↵
## Заключение↵
↵
В данной статье мы подробно рассмотрели понятие виртуального дерева, доказали его основные свойства, описали эффективный алгоритм построения и продемонстрировали применение этой техники для решения задачи на ДП в деревьях. Преимущество данного подхода состоит в том, что он позволяет свести задачу к работе с подструктурой размером $O(|X|)$, что существенно ускоряет вычисления даже при большом размере исходного дерева.↵
↵
Статья основана на [разборе](https://atcoder.jp/contests/abc340/editorial/9256) задачи с Atcoder (ABC340G) от [user:Nyaan,2025-02-26].
↵
## Введение↵
↵
Во многих задачах, связанных с деревьями, требуется работать с подмножеством вершин, сохраняя при этом структуру дерева, индуцированную их попарными наименьшими общими предками (LCA). Понятие *виртуального дерева* позволяет перейти от исходного дерева $T$ с $N$ вершинами к подструктуре, размер которой линейно зависит от размера выбранного множества $X$. Это существенно ускоряет алгоритмы, в частности, для подсчета ДП.↵
↵
## Определение виртуального дерева↵
↵
Пусть $T$ — заданное корневое дерево, а $X$ — некоторое подмножество его вершин. **Виртуальное дерево** $A(X)$ определяется следующим образом:↵
$$↵
A(X) = \{ \operatorname{lca}(x,y) \mid x,y \in X \},↵
$$↵
где $\operatorname{lca}(x,y)$ означает наименьшего общего предка вершин $x$ и $y$ в дереве $T$. В этом дереве между каждой парой вершин проводится ребро, если одна из них является предком другой в $T$.↵
↵
Такое определение гарантирует, что в структуру попадают все вершины, «важные» для анализа $X$ (то есть все LCA для пар вершин из $X$), а связность наследуется от исходного дерева.↵
↵
## Ключевые свойства и их доказательства↵
↵
### Свойство 1: Оценка по размеру↵
↵
**Утверждение.** ↵
Для любого множества $X$ вершин дерева $T$ справедливо неравенство:↵
$$↵
\vert A(X) \vert \le 2 \vert X \vert - 1.↵
$$↵
↵
**Доказательство.** ↵
Выберем порядок обхода в глубину (DFS) и рассмотрим вершины множества $X$, расположенные в порядке обхода: ↵
$$↵
x_1, x_2, \dots, x_m.↵
$$↵
Обозначим↵
$$↵
l = \operatorname{lca}(x_1, x_2).↵
$$↵
Далее необходимо доказать следующую лемму.↵
↵
> **Лемма 1.** Для любого целого числа $n \ge 3$, если $\operatorname{lca}(x_1,x_n) \neq l$, то↵
$$↵
\operatorname{lca}(x_1,x_n) = \operatorname{lca}(x_2,x_n).↵
$$↵
>↵
> **Доказательство (от противного).** ↵
> Предположим, что для некоторого $n \ge 3$ выполняется $\operatorname{lca}(x_1,x_n) \neq l$ и одновременно $\operatorname{lca}(x_1,x_n) \neq \operatorname{lca}(x_2,x_n)$. Пусть $k = \operatorname{lca}(x_1,x_n)$. ↵
> Если $k$ является предком $l$, то по определению получаем, что и $\operatorname{lca}(x_1,x_n)$, и $\operatorname{lca}(x_2,x_n)$ равны $k$, что противоречит предположению. Следовательно, $k$ должно быть потомком $l$. Но тогда в порядке DFS должен получиться порядок $x_1, x_n, x_2$, что противоречит выбранному порядку обхода. ↵
> Противоречие завершает доказательство леммы.↵
↵
Возвращаясь к доказательству свойства, заметим, что согласно лемме, каждый LCA, возникающий при последовательном рассмотрении вершин $X$, равен либо:↵
1. самой вершине $x_1$,↵
2. $\operatorname{lca}(x_1,x_2)$, либо↵
3. $\operatorname{lca}(x_2,x_n)$ для некоторого $n$.↵
↵
Таким образом, можно рекурсивно записать:↵
$$↵
A(\{x_1,\dots,x_m\}) = A(\{x_2,\dots,x_m\}) \cup \{ x_1, \operatorname{lca}(x_1,x_2) \}.↵
$$↵
Из этого следует, что↵
$$↵
\vert A(X) \vert \le \vert A(\{x_2,\dots,x_m\}) \vert + 2.↵
$$↵
Применяя индукцию по размеру множества $X$, получаем итоговую оценку:↵
$$↵
\vert A(X) \vert \le 2 \vert X \vert - 1.↵
$$↵
При этом, если дерево $T$ является совершенным двоичным деревом, а $X$ состоит из его листьев, неравенство достигается в точности.↵
↵
### Свойство 2: Представление через последовательные LCA↵
↵
**Утверждение.** ↵
Пусть вершины $X$ упорядочены по порядку обхода DFS: $x_1, x_2, \dots, x_m$. Тогда↵
$$↵
A(X) = X \cup \{ \operatorname{lca}(x_i, x_{i+1}) \mid 1 \le i < m \}.↵
$$↵
↵
**Доказательство.** ↵
Начнём с рекурсивного представления:↵
$$↵
A(\{x_1,\dots,x_m\}) = A(\{x_2,\dots,x_m\}) \cup \{ x_1, \operatorname{lca}(x_1,x_2) \}.↵
$$↵
Раскроем его рекурсивно:↵
$$↵
\begin{aligned}↵
A(\{x_1,\dots,x_m\}) &= A(\{x_2,\dots,x_m\}) \cup \{ x_1, \operatorname{lca}(x_1,x_2) \} \\↵
&= A(\{x_3,\dots,x_m\}) \cup \{ x_1, \operatorname{lca}(x_1,x_2), x_2, \operatorname{lca}(x_2,x_3) \} \\↵
&\quad \vdots \\↵
&= \{ x_1, \operatorname{lca}(x_1,x_2), x_2, \dots, x_{m-1}, \operatorname{lca}(x_{m-1},x_m), x_m \}.↵
\end{aligned}↵
$$↵
Таким образом, получаем требуемое представление.↵
↵
## Построение виртуального дерева↵
↵
Дано множество вершин $X$. Виртуальное дерево $A(X)$ можно построить за время↵
$$↵
O\left(|X| (\log |X| + \log N)\right),↵
$$↵
при этом необходимо выполнить предпосчет исходного дерева $T$ за $O(N \log N)$ для обеспечения быстрых запросов LCA.↵
↵
**Основные шаги построения:**↵
↵
1. **Предподсчет.** ↵
С помощью обхода в глубину (DFS) или методов, таких как двоичные подъемы или HLD, вычисляем времена входа в вершины и LCA для любых двух вершин.↵
↵
2. **Сортировка вершин.** ↵
Сортируем вершины множества $X$ в порядке их появления при обходе DFS (используя время входа).↵
↵
3. **Добавление LCA.** ↵
Для последовательных вершин $x_i$ и $x_{i+1}$ вычисляем $\operatorname{lca}(x_i,x_{i+1})$. Объединение $X$ и этих LCA даёт множество вершин $A(X)$ (согласно свойству 2).↵
↵
4. **Построение дерева.** ↵
Получив множество вершин $A(X)$, отсортированных по порядку DFS, можно пройти по последовательности, используя стек, для восстановления структуры дерева. При обходе поддерживается стек, верхний элемент которого является текущей вершиной, а если новая вершина не является потомком вершины на вершине стека, производится «выталкивание» элементов до нахождения соответствующего предка, после чего производится соединение.↵
↵
Такой алгоритм гарантирует, что виртуальное дерево содержит $O(|X|)$ вершин, что позволяет эффективно применять ДП.↵
↵
## Применение: подсчет поддеревьев в раскрашенном дереве↵
↵
### Условие [задачи](https://atcoder.jp/contests/abc340/tasks/abc340_g)↵
↵
Дано дерево $T$ с $N$ вершинами, пронумерованными от $1$ до $N$. Ребро $i$ соединяет вершины $u[i]$ и $v[i]$. Каждая вершина $i$ раскрашена в цвет $A[i]$. ↵
Необходимо найти (по модулю 998244353) количество (непустых) подмножеств $S$ вершин дерева $T$, удовлетворяющих условию:↵
- Индуцированный граф $G[S]$ является деревом.↵
- Все вершины $G[S]$ со степенью 1 имеют один и тот же цвет.↵
↵
### Основная идея решения↵
↵
Основная идея заключается в разбиении задачи по цветам. Для каждого цвета $c$ рассматривается множество вершин $X$, раскрашенных в $c$, и строится виртуальное дерево для $X$. Затем на полученном дереве выполняется ДП для подсчета допустимых поддеревьев, у которых все вершины со степенью 1 имеют цвет $c$. Итоговый ответ получается суммированием результатов по всем цветам.↵
↵
Благодаря построению виртуального дерева, хотя исходное дерево содержит $N$ вершин, каждое виртуальное дерево для конкретного цвета имеет размер $O(|X|)$, что позволяет выполнять DP за приемлемое время.↵
↵
Получаем, что суммарная сложность предпосчета и построения виртуальных деревьев не превышает $O(N \log N + \sum_{c \in C} |X_c| (\log |X_c| + \log N)) = O(N \log N)$, где $C$ это множество различных цветов вершин и $X_c$ это множество вершин с цветом $c$.↵
↵
### Реализация решения↵
↵
Ниже приведён полный код на C++ с комментариями, описывающими основные компоненты алгоритма:↵
↵
```cpp↵
#include <bits/stdc++.h>↵
using namespace std;↵
↵
const int MAXN = 2e5 + 10;↵
const int LOGN = 20;↵
const int MOD = 998244353;↵
↵
vector<int> g[MAXN]; // Список смежности заданного графа↵
vector<int> ver[MAXN]; // Для каждого цвета ver[c] хранит вершины с цветом c↵
vector<int> aux_g[MAXN]; // Список смежности для виртуального дерева↵
int col[MAXN]; // Цвет каждой вершины↵
int tmr, n; // Глобальный счетчик времени и количество вершин↵
int up[LOGN][MAXN], dep[MAXN], tin[MAXN]; // Для подсчета LCA и времени входа↵
↵
// DP массивы для динамического программирования на виртуальном дереве↵
int dp[MAXN][2], sum[MAXN];↵
↵
// Функция предподсчета: DFS для вычисления tin, массива up и dep↵
void dfs_precalc(int v, int p) {↵
tin[v] = ++tmr;↵
up[0][v] = p;↵
dep[v] = dep[p] + 1;↵
for (int i = 1; i < LOGN; ++i)↵
up[i][v] = up[i - 1][up[i - 1][v]];↵
for (auto to : g[v])↵
if (to != p)↵
dfs_precalc(to, v);↵
}↵
↵
// Функция для вычисления LCA с использованием двоичных подъемов↵
int getlca(int x, int y) {↵
if (dep[x] < dep[y]) swap(x, y);↵
for (int i = LOGN - 1; i >= 0; --i)↵
if (dep[up[i][x]] >= dep[y])↵
x = up[i][x];↵
if (x == y) return x;↵
for (int i = LOGN - 1; i >= 0; --i)↵
if (up[i][x] != up[i][y]) {↵
x = up[i][x];↵
y = up[i][y];↵
}↵
return up[0][x];↵
}↵
↵
// DFS по виртуальному дереву для выполнения DP.↵
// Параметр c — целевой цвет, для которого производится подсчет.↵
↵
void dfs_calc(int v, int p, int c, int &ans) {↵
dp[v][0] = dp[v][1] = 0;↵
sum[v] = 0;↵
for(auto to : aux_g[v]) {↵
if(to == p) continue;↵
dfs_calc(to, v, c, ans);↵
// Переходы DP: объединение текущего состояния с результатом из поддерева.↵
int nxt0 = (dp[v][0] + sum[to]) % MOD;↵
int nxt1 = ((dp[v][0] + dp[v][1]) * 1ll * sum[to] % MOD + dp[v][1]) % MOD;↵
dp[v][0] = nxt0;↵
dp[v][1] = nxt1;↵
}↵
sum[v] = (dp[v][0] + dp[v][1]) % MOD;↵
if(col[v] == c) {↵
// Если вершина имеет целевой цвет, она может участвовать в корректном поддереве.↵
sum[v] = (sum[v] + 1) % MOD;↵
ans = (ans + sum[v]) % MOD;↵
} else {↵
ans = (ans + dp[v][1]) % MOD;↵
}↵
}↵
↵
// Функция для построения виртуального дерева для цвета c и выполнения DP.↵
void calc_aux(int c, int &ans) {↵
auto p = ver[c];↵
if (p.empty()) return;↵
// Сортируем вершины по времени входа (tin) — порядку обхода DFS.↵
sort(p.begin(), p.end(), [&](const int a, const int b) { return tin[a] < tin[b]; });↵
vector<int> verstk = {1}; // Инициализируем стек с корнем дерева T (вершина 1).↵
aux_g[1].clear();↵
auto add = [&](int u, int v) {↵
aux_g[u].push_back(v);↵
aux_g[v].push_back(u);↵
};↵
// Обрабатываем каждую вершину из множества p, поддерживая стек для построения виртуального дерева.↵
for (auto u : p) {↵
if (u == 1) continue;↵
int lca = getlca(u, verstk.back());↵
if (lca != verstk.back()) {↵
while (verstk.size() >= 2 && tin[lca] < tin[verstk[verstk.size() - 2]]) {↵
add(verstk.back(), verstk[verstk.size() - 2]);↵
verstk.pop_back();↵
}↵
if (verstk.size() >= 2 && tin[lca] != tin[verstk[verstk.size() - 2]]) {↵
aux_g[lca].clear();↵
add(verstk.back(), lca);↵
verstk.back() = lca;↵
} else {↵
add(verstk.back(), lca);↵
verstk.pop_back();↵
}↵
}↵
aux_g[u].clear();↵
verstk.push_back(u);↵
}↵
while (verstk.size() > 1) {↵
add(verstk.back(), verstk[verstk.size() - 2]);↵
verstk.pop_back();↵
}↵
// Выполняем DP по виртуальному дереву, начиная с корня 1.↵
return dfs_calc(1, 0, c, ans);↵
}↵
↵
int main() {↵
ios_base::sync_with_stdio(0);↵
cin.tie(0);↵
cin >> n;↵
// Считываем цвета и группируем вершины по цвету.↵
for (int i = 1; i <= n; ++i) {↵
cin >> col[i];↵
ver[col[i]].push_back(i);↵
}↵
for (int i = 1; i < n; ++i) {↵
int u, v;↵
cin >> u >> v;↵
g[u].push_back(v);↵
g[v].push_back(u);↵
}↵
dfs_precalc(1, 0);↵
int ans = 0;↵
// Для каждого возможного цвета обрабатываем соответствующее виртуальное дерево.↵
for (int i = 1; i <= n; ++i)↵
calc_aux(i, ans);↵
cout << ans;↵
return 0;↵
}↵
```↵
↵
### Объяснение реализации↵
↵
- **Предподсчет.** ↵
Функция `dfs_precalc` выполняет обход дерева $T$ в глубину, вычисляя время входа (`tin`), глубину и заполняя таблицу двоичных подъемов `up` для быстрых запросов LCA.↵
↵
- **Вычисление LCA.** ↵
Функция `getlca` реализует двоичные подъемы для быстрого нахождения наименьшего общего предка двух вершин.↵
↵
- **Построение виртуального дерева.** ↵
Функция `calc_aux` для каждого цвета $c$ берёт множество вершин `ver[c]`, сортирует его по `tin` и с помощью стека строит виртуальное дерево. При этом для каждой пары последовательных вершин вычисляется LCA, что соответствует свойству 2.↵
↵
- **Динамическое программирование.** ↵
Функция `dfs_calc` выполняет обход виртуального дерева и объединяет результаты поддеревьев согласно переходам DP. Состояния DP рассчитаны таким образом, что учитывается вклад вершины, если она имеет целевой цвет, и производится подсчет только тех поддеревьев, где все листья имеют один цвет.↵
↵
- **Сбор результата.** ↵
Функция `main` считывает входные данные, выполняет предподсчет, и затем суммирует результаты для каждого цвета, выводя итоговый ответ по модулю 998244353.↵
↵
## Заключение↵
↵
В данной статье мы подробно рассмотрели понятие виртуального дерева, доказали его основные свойства, описали эффективный алгоритм построения и продемонстрировали применение этой техники для решения задачи на ДП в деревьях. Преимущество данного подхода состоит в том, что он позволяет свести задачу к работе с подструктурой размером $O(|X|)$, что существенно ускоряет вычисления даже при большом размере исходного дерева.↵
↵
Статья основана на [разборе](https://atcoder.jp/contests/abc340/editorial/9256) задачи с Atcoder (ABC340G) от [user:Nyaan,2025-02-26].
