ДП на дереве: как посчитать размеры всех поддеревьев одним DFS?
Дано корневое дерево на $n$ вершинах списком смежности, корень — вершина 1. Нужно для каждой вершины узнать размер её поддерева (число вершин в нём, включая саму). Понимаю, что это базовое ДП на дереве, но хочу увидеть аккуратную реализацию без рекурсии в глубину $10^5$, чтобы не было переполнения стека.
2 ответа
Размер поддерева — каноничное ДП на дереве: sz[v] = 1 + сумма sz[u] по детям u. Считается одним post-order DFS: сначала спускаемся в детей, потом агрегируем.
Рекурсивно (просто, но осторожно со стеком при глубине ~$10^5$):
int n;
vector<vector<int>> g(n + 1);
vector<int> sz(n + 1, 0);
void dfs(int v, int parent) {
sz[v] = 1;
for (int u : g[v])
if (u != parent) {
dfs(u, v);
sz[v] += sz[u]; // агрегация снизу вверх
}
}
// вызов: dfs(1, 0);
Сложность $O(n)$ времени и памяти — каждое ребро проходим один раз. Это база для большинства задач на деревьях: число вершин на чётной/нечётной глубине, сумма расстояний, центроидная декомпозиция и т.д.
Про переполнение стека — реальная проблема на Codeforces при глубине дерева до $10^5$ (бамбук). Варианты:
- Итеративный DFS через явный стек. Для размеров поддеревьев удобно сделать топологический порядок: получить порядок вершин в обходе, затем идти в обратном порядке и прибавлять
sz[v]к родителю. - Увеличить стек: запустить основную логику в отдельном потоке с увеличенным размером, либо
#pragma comment(linker, "/STACK:268435456")(MSVC).
Итеративный вариант через обратный порядок:
vector<int> order; // заполняется обычным DFS/BFS в порядке посещения
for (int idx = order.size() - 1; idx >= 0; --idx) {
int v = order[idx];
sz[v] = 1;
for (int u : g[v]) if (u != par[v]) sz[v] += sz[u];
}