Root-ты ауыстыру деген не?

Root-ты ауыстыру, немесе rerooting, — Tree DP әдісі. Ол answer бір root үшін ғана емес, әр вершина root болған жағдай үшін керек кезде қолданылады.

Неге керек?

Егер әр вершинадан бөлек DFS жасасақ, өте баяу болады.

Мысалы, n = 100000 болса:

Сондықтан rerooting көбіне answer барлық вершина үшін керек болғанда қолданылады.

Идея қарапайым тілмен

Бізде v вершинасы үшін answer белгілі болсын. Енді оның child-ы to үшін answer тапқымыз келеді.

Root ауысқанда ағаш былай өзгереді:

Біз answer қалай өзгеретінін түсінуіміз керек. Көп задачада оны O(1) формуламен жасауға болады.

Классикалық задача

n вершинадан тұратын ағаш берілген. Әр вершина v үшін v-дан барлық басқа вершинаға дейінгі distance суммасын табу керек.

Яғни:

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 вершинасының поддерево size-ы
• dp[v] — v-дан өзінің поддеревосындағы вершинаға дейінгі distance суммасы
• ans[v] — v-дан барлық вершинаға дейінгі distance суммасы

Бірінші DFS

Алдымен ағашты root = 1 қылып аламыз. Бірінші DFS ішінде санаймыз:

• sz[v]
• dp[v]

Әр вершина үшін бастапқы мән:

sz[v] = 1;

Кейін child-тардан қосамыз:

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] — root 1-ден барлық вершинаға дейінгі distance суммасы.

Rerooting-тің негізгі формуласы

Root-ты v-дан оның child-ы to-ға ауыстырамыз.

Сонда:

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

Бұл ең маңызды формула.

Неге формула осындай?

Root 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 үшін answer білеміз. Сол арқылы to үшін answer аламыз.

Толық реализация

#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 мысалы

Input

5
1 2
1 3
3 4
3 5

Output

6 9 5 8 8

Мұнда:

ans[3]-ті тексерейік

3 вершинасынан distance-тар:

Сумма:

1 + 2 + 0 + 1 + 1 = 5

Сондықтан:

ans[3] = 5

Формула мысалда қалай жұмыс істейді?

Root-ты 1-ден 3-ке ауыстырамыз.

3 вершинасының поддеревосы:

Формула:

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

Қоямыз:

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

Rerooting жалпы схемасы

Көбіне 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 тек distance sum үшін ғана емес. Мысалы:

• әр вершинадан максималды distance
• әр вершина root болса, subtree answer
• path саны
• поддерево және поддерево сыртындағы мәндер
• барлық вершина үшін answer

Задача 2: әр вершинадан ең алыс distance

Әр вершина үшін ең алыс вершинаға дейінгі distance табу керек.

Яғни:

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

Мұнда сақтаймыз:

• down[v] — v-дан төмен қарай ең үлкен depth
• up[v] — v-дан жоғары немесе басқа поддеревоға кететін ең жақсы path

down[v]

down[v] обычный DFS арқылы саналады:

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

Егер child жоқ болса:

down[v] = 0;

up[to]

up[to] санау үшін v-дан шығатын, бірақ to поддеревосына кірмейтін ең жақсы жолды білу керек.

Сондықтан әр вершина үшін child-тардан келетін екі max depth сақтаймыз:

best1
best2

Егер child to best1 берген болса, онда оған best2 қолданамыз. Әйтпесе best1 қолданамыз.

Код: әр вершинадан максималды distance

#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 қашан қолданамыз?

Задачада мына сөздер болса, rerooting керек болуы мүмкін:

• әр вершина үшін
• егер осы вершина root болса
• барлық вершина үшін answer
• әр вершинадан distance суммасы
• әр вершинадан максималды distance
• answer поддеревоға және поддерево сыртындағы бөлікке байланысты

Типтік қателер

• sz[v] санауды ұмыту
• sz[to] және n - sz[to] шатастыру
• тек бірінші DFS жасап қою
• ans[1] = dp[1] қоюды ұмыту
• parent-қа қайта кіріп кету
• қай вершина жақындап, қай вершина алыстағанын түсінбей формула жазу
• rerooting әрқашан бір формула деп ойлау

Асимптотика

Distance sum қысқа шаблон

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);
    }
}

Ең бастысы