Что такое перевешивание корня?

Перевешивание корня, или rerooting, — это метод DP на дереве. Он нужен, когда надо посчитать ответ не только для одного корня, а для каждой вершины как для корня.

Зачем это нужно?

Если для каждой вершины запускать отдельный DFS, будет слишком медленно.

Например, если n = 100000:

Поэтому перевешивание корня используется в задачах, где ответ нужен для всех вершин дерева.

Идея на словах

Допустим, мы знаем ответ для вершины v. Теперь хотим сделать корнем ее ребенка to.

Тогда дерево как будто перевешивается:

Мы должны понять, как изменится ответ. Часто это можно сделать формулой за O(1).

Классическая задача

Дано дерево из n вершин. Для каждой вершины v нужно найти сумму расстояний от v до всех остальных вершин.

То есть надо посчитать:

ans[v] = dist(v, 1) + dist(v, 2) + ... + dist(v, n)

Для всех v.

Пример

Дерево:

Если взять вершину 1:

Значит:

ans[1] = 6

Что будем хранить?

Нам нужны массивы:

sz[v]
dp[v]
ans[v]

• sz[v] — размер поддерева вершины v
• dp[v] — сумма расстояний от v до вершин в ее поддереве
• ans[v] — сумма расстояний от v до всех вершин дерева

Первый DFS

Сначала подвесим дерево за вершину 1. В первом DFS считаем:

• sz[v]
• dp[v]

Для каждой вершины:

sz[v] = 1;

Потом добавляем детей:

sz[v] += sz[to];
dp[v] += dp[to] + sz[to];

Код первого DFS

void dfs1(int v, int p) {
    sz[v] = 1;
    dp[v] = 0;

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs1(to, v);

        sz[v] += sz[to];
        dp[v] += dp[to] + sz[to];
    }
}

После этого:

ans[1] = dp[1];

Потому что dp[1] — сумма расстояний от корня 1 до всех вершин.

Главная формула перевешивания

Пусть мы переносим корень из v в его ребенка to.

Тогда:

ans[to] = ans[v] - sz[to] + (n - sz[to]);

Это самая важная формула.

Почему формула такая?

Когда корень переходит из v в to:

• все вершины внутри поддерева to становятся ближе на 1
• таких вершин sz[to]
• поэтому сумма уменьшается на sz[to]

• все остальные вершины становятся дальше на 1
• таких вершин n - sz[to]
• поэтому сумма увеличивается на n - sz[to]

Второй DFS

Во втором DFS уже считаем ans для всех вершин.

void dfs2(int v, int p) {
    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        ans[to] = ans[v] - sz[to] + (n - sz[to]);

        dfs2(to, v);
    }
}

Мы знаем ответ для v и через него получаем ответ для to.

Полная реализация

#include <bits/stdc++.h>
using namespace std;

const int N = 100010;

int n;
vector<int> g[N];
long long sz[N];
long long dp[N];
long long ans[N];

void dfs1(int v, int p) {
    sz[v] = 1;
    dp[v] = 0;

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs1(to, v);

        sz[v] += sz[to];
        dp[v] += dp[to] + sz[to];
    }
}

void dfs2(int v, int p) {
    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        ans[to] = ans[v] - sz[to] + (n - sz[to]);

        dfs2(to, v);
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;

    for (int i = 1; i <= n - 1; i++) {
        int a, b;
        cin >> a >> b;

        g[a].push_back(b);
        g[b].push_back(a);
    }

    dfs1(1, 0);

    ans[1] = dp[1];

    dfs2(1, 0);

    for (int i = 1; i <= n; i++) {
        cout << ans[i] << ' ';
    }

    return 0;
}

Пример входа

Input

5
1 2
1 3
3 4
3 5

Output

6 9 5 8 8

Здесь:

Проверим ans[3]

Расстояния от вершины 3:

Сумма:

1 + 2 + 0 + 1 + 1 = 5

Поэтому:

ans[3] = 5

Как формула работает на примере?

Переносим корень из 1 в 3.

Поддерево вершины 3:

Формула:

ans[3] = ans[1] - sz[3] + (n - sz[3]);

Подставим:

ans[3] = 6 - 3 + (5 - 3) = 5

Общая схема rerooting

Обычно перевешивание корня делается так:

Шаблон rerooting

void dfs1(int v, int p) {
    // считаем dp[v] снизу вверх

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs1(to, v);

        // dp[v] обновляется через dp[to]
    }
}

void dfs2(int v, int p) {
    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        // ans[to] пересчитывается через ans[v]

        dfs2(to, v);
    }
}

Еще одна задача: количество вершин на расстоянии

Иногда rerooting используется не только для сумм расстояний. Например, можно считать:

• максимальное расстояние от каждой вершины
• количество путей
• значения в поддеревьях и вне поддеревьев
• ответ для каждой вершины как для корня

Задача 2: максимальное расстояние от каждой вершины

Для каждой вершины нужно найти расстояние до самой далекой вершины.

Например:

ans[v] = max distance from v to any other vertex

Тут удобно хранить:

• down[v] — максимальная глубина вниз в поддереве v
• up[v] — лучший путь из v наверх или в другие поддеревья

down[v]

down[v] считается обычным DFS:

down[v] = max(down[to] + 1);

Если у вершины нет детей:

down[v] = 0;

up[to]

Чтобы посчитать up[to], нужно знать лучший путь из v, который не идет в поддерево to.

Поэтому для каждой вершины храним две максимальные глубины среди детей:

best1
best2

Если ребенок to дал best1, то для него используем best2. Иначе используем best1.

Код: максимальное расстояние для каждой вершины

#include <bits/stdc++.h>
using namespace std;

const int N = 100010;

int n;
vector<int> g[N];
int downDist[N];
int upDist[N];
int ans[N];

void dfs1(int v, int p) {
    downDist[v] = 0;

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs1(to, v);

        downDist[v] = max(downDist[v], downDist[to] + 1);
    }
}

void dfs2(int v, int p) {
    int best1 = -1;
    int best2 = -1;

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        int cur = downDist[to] + 1;

        if (cur > best1) {
            best2 = best1;
            best1 = cur;
        } else if (cur > best2) {
            best2 = cur;
        }
    }

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        int use = best1;

        if (downDist[to] + 1 == best1) {
            use = best2;
        }

        upDist[to] = upDist[v] + 1;

        if (use != -1) {
            upDist[to] = max(upDist[to], use + 1);
        }

        dfs2(to, v);
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;

    for (int i = 1; i <= n - 1; i++) {
        int a, b;
        cin >> a >> b;

        g[a].push_back(b);
        g[b].push_back(a);
    }

    dfs1(1, 0);

    upDist[1] = 0;

    dfs2(1, 0);

    for (int i = 1; i <= n; i++) {
        ans[i] = max(downDist[i], upDist[i]);
        cout << ans[i] << ' ';
    }

    return 0;
}

Когда использовать перевешивание корня?

Используй rerooting, если в задаче есть слова:

• для каждой вершины
• если выбрать вершину корнем
• ответ для всех вершин дерева
• сумма расстояний от каждой вершины
• максимальное расстояние от каждой вершины
• ответ зависит от поддерева и остальной части дерева

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

• Забыть посчитать sz[v]
• Перепутать sz[to] и n - sz[to]
• Запустить только первый DFS
• Забыть поставить ans[1] = dp[1]
• Зайти обратно в родителя
• Писать формулу без понимания, какие вершины стали ближе и дальше
• Думать, что rerooting всегда одна и та же формула

Асимптотика

Короткий шаблон суммы расстояний

void dfs1(int v, int p) {
    sz[v] = 1;

    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        dfs1(to, v);

        sz[v] += sz[to];
        dp[v] += dp[to] + sz[to];
    }
}

void dfs2(int v, int p) {
    for (int to : g[v]) {
        if (to == p) {
            continue;
        }

        ans[to] = ans[v] - sz[to] + (n - sz[to]);

        dfs2(to, v);
    }
}

Главное запомнить