Rerooting (смена корня): как за O(n) найти для каждой вершины сумму расстояний до всех остальных?
Дерево на вершинах. Нужно для каждой вершины посчитать — сумму расстояний от неё до всех остальных. Наивно (BFS из каждой) — TLE. Знаю, что есть техника rerooting, но не понимаю, как пересчитывать ответ при «переносе корня» на соседа.
2 ответа
Это эталонная задача на rerooting (смену корня) в два прохода DFS.
Проход 1 (снизу вверх): считаем sz[v] (размер поддерева) и ans[root] — сумму расстояний от корня. down[v] = сумма расстояний от v до вершин его поддерева, считается как .
Проход 2 (сверху вниз, rerooting): зная ответ для родителя p, получаем ответ для ребёнка v по формуле переноса корня. Когда корень переезжает с p на v: вершины поддерева v стали ближе на 1 (их sz[v] штук), все остальные — дальше на 1 (n - sz[v] штук):
int n;
vector<vector<int>> g;
vector<long long> sz, ans;
void dfs1(int v, int p) {
sz[v] = 1; ans[v] = 0;
for (int u : g[v]) if (u != p) {
dfs1(u, v);
sz[v] += sz[u];
ans[v] += ans[u] + sz[u]; // сумма расст. внутри поддерева v
}
}
void dfs2(int v, int p) { // ans[v] уже = глобальный ответ
for (int u : g[v]) if (u != p) {
ans[u] = ans[v] - sz[u] + (n - sz[u]); // перенос корня v -> u
dfs2(u, v);
}
}
// dfs1(0, -1); ans[0] уже глобальный после dfs1
// dfs2(0, -1);
Сложность времени и памяти. long long обязателен: сумма расстояний до по величине.
Тонкость, на которой все спотыкаются: после dfs1 ans[root] уже корректен (это сумма расстояний от корня), а вот ans[v] для остальных в dfs1 — это лишь сумма по поддереву, не глобальный ответ. Глобальные значения проставляет именно dfs2 через формулу переноса, и в момент входа в dfs2(v,...) значение ans[v] уже глобальное (его записал родитель). Порядок строго: сначала пересчитать ans[u] из ans[v], потом рекурсия в u.
Общий шаблон rerooting шире: так считают «максимальную глубину из каждой вершины», диаметр-производные величины и любые агрегаты, где переход корня к соседу меняет вклад предсказуемо.