Rerooting DP

Dời gốc

Chọn gốc 1. Gọi $sz_u$ là kích thước cây con của $u$ và $f(u)$ là tổng khoảng cách từ $u$ tới mọi đỉnh. Lượt dưới lên cho $f(1)$, là tổng độ sâu của mọi đỉnh.

Bây giờ dời gốc từ $u$ sang con $v$. $sz_v$ đỉnh trong cây con của $v$ gần thêm 1, còn $n - sz_v$ đỉnh còn lại xa thêm 1. Vậy

Đi từ gốc xuống theo thứ tự BFS (hoặc thứ tự $tin$), mỗi đỉnh tính từ cha trong $O(1)$.

Để tránh đệ quy sâu khi cây là một đường dài $2 \cdot 10^5$, lấy thứ tự BFS từ gốc: duyệt ngược thứ tự đó là đi từ lá lên (con xong trước cha), duyệt xuôi là đi từ gốc xuống (cha xong trước con).

Ví dụTổng Khoảng Cách Từ Mỗi Điểm

Đề. Cây $n \le 2 \cdot 10^5$ đỉnh. Với mỗi đỉnh, in tổng khoảng cách tới mọi đỉnh.

Lời giải. Với gốc 1, tổng khoảng cách $f(1)$ là tổng độ sâu mọi đỉnh. Dời gốc từ $u$ sang con $v$: $sz_v$ đỉnh trong cây con của $v$ gần thêm 1, $n - sz_v$ đỉnh còn lại xa thêm 1, nên $f(v) = f(u) - sz_v + (n - sz_v)$. Cả bài không cần đệ quy: một thứ tự BFS cho phép đi “từ lá lên” bằng cách duyệt ngược, và “từ gốc xuống” bằng cách duyệt xuôi.

Bước 1: BFS lấy thứ tự. ord xếp cha trước con, cùng par và dep. Vòng for trên ord vừa duyệt vừa thêm, chính là BFS dùng ord làm hàng đợi.

ord.push_back(1);
par[1] = 0;
for (ll i = 0; i < (ll)ord.size(); i++) {
    ll u = ord[i];
    for (ll v : g[u])
        if (v != par[u]) {
            par[v] = u;
            dep[v] = dep[u] + 1;
            ord.push_back(v);
        }
}

Bước 2: từ lá lên. Duyệt ord ngược: con xong trước cha, cộng sz lên cha. Cùng vòng đó cộng mọi dep vào f[1].

for (ll i = n - 1; i >= 0; i--) {
    ll u = ord[i];
    sz[u]++;
    if (par[u]) sz[par[u]] += sz[u];
    f[1] += dep[u];
}

Bước 3: dời gốc. Duyệt xuôi từ phần tử thứ hai; cha luôn có f trước con.

for (ll i = 1; i < n; i++) {
    ll v = ord[i];
    f[v] = f[par[v]] - sz[v] + (n - sz[v]);
}

Tổng tới khoảng $2 \cdot 10^{10}$, cần ll.

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

const ll N = 2e5 + 5;
ll n, par[N], dep[N], sz[N], f[N];
vector<ll> g[N], ord;

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    cin >> n;
    for (ll i = 1; i < n; i++) {
        ll u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    ord.push_back(1);
    par[1] = 0;
    for (ll i = 0; i < (ll)ord.size(); i++) {
        ll u = ord[i];
        for (ll v : g[u])
            if (v != par[u]) {
                par[v] = u;
                dep[v] = dep[u] + 1;
                ord.push_back(v);
            }
    }
    for (ll i = n - 1; i >= 0; i--) {
        ll u = ord[i];
        sz[u]++;
        if (par[u]) sz[par[u]] += sz[u];
        f[1] += dep[u];
    }
    for (ll i = 1; i < n; i++) {
        ll v = ord[i];
        f[v] = f[par[v]] - sz[v] + (n - sz[v]);
    }
    for (ll u = 1; u <= n; u++) cout << f[u] << ' ';
    cout << '\n';
}

ChạyInput5
1 2
2 3
2 4
4 5

Output8 5 8 6 9

Ba vòng tuyến tính: $O(n)$.

Cạnh có hướng

Khi mỗi cạnh có một chi phí phụ thuộc vào chiều nhìn từ gốc, dời gốc qua cạnh $(u, v)$ chỉ đổi chiều nhìn của đúng cạnh đó; mọi cạnh khác vẫn được nhìn từ phía cũ. Đáp án mới bằng đáp án cũ trừ chi phí cạnh theo chiều cũ, cộng chi phí theo chiều mới. Lưu mỗi cạnh hai lần trong danh sách kề kèm chi phí theo chiều duyệt là đủ.

Ví dụTrạm Phát Trung Tâm

Đề. Cây $n \le 2 \cdot 10^5$ đỉnh với các cạnh có hướng. Chọn một trạm làm trung tâm và đảo ít cạnh nhất để từ đó đi được tới mọi trạm. In số cạnh ít nhất và mọi trạm đạt được, tăng dần.

Lời giải. Với gốc cố định, cạnh nào trỏ ngược chiều từ gốc ra (từ con về cha) thì phải đảo, nên đáp án là số cạnh như vậy. Dời gốc từ $u$ sang con $v$ chỉ làm đúng cạnh $(u, v)$ đổi chiều nhìn: chi phí của nó đổi từ $c$ thành $1 - c$, mọi cạnh khác giữ nguyên.

Bước 1: lưu cạnh hai lần. Cạnh $u \to v$ lưu thành (v, 0) ở $u$ (đi đúng chiều, không tốn) và (u, 1) ở $v$ (đi ngược, phải đảo).

g[u].push_back({v, 0});
g[v].push_back({u, 1});

Bước 2: BFS từ gốc 1. pc[v] là chi phí của cạnh nối $v$ với cha theo chiều cha xuống con; cộng chúng vào f[1].

for (pair<ll, ll> e : g[u])
    if (e.first != par[u]) {
        par[e.first] = u;
        pc[e.first] = e.second;
        f[1] += e.second;
        ord.push_back(e.first);
    }

Bước 3: dời gốc.

for (ll i = 1; i < n; i++) {
    ll v = ord[i];
    f[v] = f[par[v]] - pc[v] + (1 - pc[v]);
}

In giá trị nhỏ nhất rồi mọi đỉnh đạt nó theo thứ tự tăng.

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

const ll N = 2e5 + 5;
ll n, par[N], pc[N], f[N];
vector<pair<ll, ll>> g[N];
vector<ll> ord;

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    cin >> n;
    for (ll i = 1; i < n; i++) {
        ll u, v;
        cin >> u >> v;
        g[u].push_back({v, 0});
        g[v].push_back({u, 1});
    }
    ord.push_back(1);
    for (ll i = 0; i < (ll)ord.size(); i++) {
        ll u = ord[i];
        for (pair<ll, ll> e : g[u])
            if (e.first != par[u]) {
                par[e.first] = u;
                pc[e.first] = e.second;
                f[1] += e.second;
                ord.push_back(e.first);
            }
    }
    for (ll i = 1; i < n; i++) {
        ll v = ord[i];
        f[v] = f[par[v]] - pc[v] + (1 - pc[v]);
    }
    ll best = *min_element(f + 1, f + n + 1);
    cout << best << '\n';
    for (ll u = 1; u <= n; u++)
        if (f[u] == best) cout << u << ' ';
    cout << '\n';
}

ChạyInput5
2 1
2 3
4 2
4 5

Output0
4

$O(n)$.

Phần phía trên

Không phải đại lượng nào cũng dời được bằng cộng trừ. Với khoảng cách xa nhất, đáp án của $v$ là max của hai phần:

• Phía dưới $down_v$: đường xa nhất đi vào cây con của $v$.
• Phía trên $up_v$: đường xa nhất đi qua cha $u$, tức $up_v = 1 + \max(up_u, \text{đường xa nhất từ } u \text{ xuống một con khác } v)$.

Phép max không trừ ngược được, nên ở mỗi đỉnh giữ hai giá trị lớn nhất của các con, cùng con cho giá trị lớn nhất. Nếu $v$ chính là con tốt nhất thì dùng giá trị thứ hai, không thì dùng giá trị thứ nhất.

Khuôn chung: $down$ tính từ lá lên, $up$ tính từ gốc xuống, và phần “mọi con trừ một” lấy từ tóm tắt các con (hai giá trị lớn nhất cho max, tổng trừ đi cho tổng, tiền tố và hậu tố cho phép gộp tùy ý).

Ví dụĐiểm Xa Nhất Trên Cây

Đề. Cây $n \le 2 \cdot 10^5$ đỉnh. Với mỗi đỉnh, in khoảng cách tới đỉnh xa nhất.

Lời giải. Đỉnh xa nhất từ $v$ hoặc nằm dưới $v$ (trong cây con), hoặc nằm phía trên (đi lên cha rồi đi đâu đó). Phần dưới tính từ lá lên; phần trên tính từ gốc xuống. Phần trên của $v$ là lên cha 1 bước, rồi hoặc đi tiếp lên trên cha, hoặc xuống một nhánh khác của cha. Nếu nhánh tốt nhất của cha chính là $v$ thì phải lấy nhánh tốt nhì, nên mỗi đỉnh giữ hai giá trị lớn nhất.

Bước 1: hai nhánh dài nhất. d1[u], d2[u] là hai giá trị lớn nhất của $1 + d1_c$ trên các con $c$ (0 nếu thiếu), bc[u] là con cho d1[u]. Duyệt ngược ord, mỗi con đưa ứng viên lên cha.

for (ll i = n - 1; i >= 1; i--) {
    ll v = ord[i], u = par[v], x = d1[v] + 1;
    if (x > d1[u]) {
        d2[u] = d1[u];
        d1[u] = x;
        bc[u] = v;
    } else if (x > d2[u]) d2[u] = x;
}

Bước 2: phần phía trên. $up_1 = 0$; duyệt xuôi, $up_v = 1 + \max(up_u, \text{nhánh tốt nhất của } u \text{ không qua } v)$.

for (ll i = 1; i < n; i++) {
    ll v = ord[i], u = par[v];
    up[v] = 1 + max(up[u], bc[u] == v ? d2[u] : d1[u]);
}

Đáp án của $v$ là $\max(d1_v, up_v)$.

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

const ll N = 2e5 + 5;
ll n, par[N], d1[N], d2[N], bc[N], up[N];
vector<ll> g[N], ord;

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    cin >> n;
    for (ll i = 1; i < n; i++) {
        ll u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    ord.push_back(1);
    for (ll i = 0; i < (ll)ord.size(); i++) {
        ll u = ord[i];
        for (ll v : g[u])
            if (v != par[u]) {
                par[v] = u;
                ord.push_back(v);
            }
    }
    for (ll i = n - 1; i >= 1; i--) {
        ll v = ord[i], u = par[v], x = d1[v] + 1;
        if (x > d1[u]) {
            d2[u] = d1[u];
            d1[u] = x;
            bc[u] = v;
        } else if (x > d2[u]) d2[u] = x;
    }
    for (ll i = 1; i < n; i++) {
        ll v = ord[i], u = par[v];
        up[v] = 1 + max(up[u], bc[u] == v ? d2[u] : d1[u]);
    }
    for (ll u = 1; u <= n; u++) cout << max(d1[u], up[u]) << ' ';
    cout << '\n';
}

ChạyInput5
1 2
2 3
2 4
4 5

Output3 2 3 2 3

$O(n)$.

Ví dụMiền Kết Nối Tốt Nhất

Đề. Cây $n \le 2 \cdot 10^5$ đỉnh trắng hoặc đen. Với mỗi đỉnh $u$, tìm max của (số trắng − số đen) trên các tập liên thông chứa $u$.

Lời giải. Gán $val = +1$ cho trắng, $-1$ cho đen; bài thành tìm tập liên thông chứa $u$ có tổng $val$ lớn nhất. Từ lá lên: $dp_u = val_u + \sum \max(0, dp_c)$ là tập tốt nhất chứa $u$ trong cây con của nó, chỉ lấy nhánh con khi nhánh đó có lợi. Từ gốc xuống: tập chứa $v$ còn có thể kéo thêm phần phía cha. Giá trị phần đó là đáp án của cha bỏ đi phần mà $v$ đã góp, tức $ans_u - \max(0, dp_v)$; phép gộp là tổng nên bỏ đóng góp của một con chỉ là phép trừ.

Bước 1: từ lá lên. dp khởi tạo bằng $\pm 1$ ngay khi đọc; duyệt ngược ord, cộng phần có lợi lên cha.

for (ll i = n - 1; i >= 1; i--) {
    ll v = ord[i];
    dp[par[v]] += max(0LL, dp[v]);
}

Bước 2: từ gốc xuống. Đáp án của gốc chính là dp[1]; con $v$ lấy thêm phần phía cha nếu nó dương.

ans[1] = dp[1];
for (ll i = 1; i < n; i++) {
    ll v = ord[i];
    ans[v] = dp[v] + max(0LL, ans[par[v]] - max(0LL, dp[v]));
}

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

const ll N = 2e5 + 5;
ll n, par[N], dp[N], ans[N];
vector<ll> g[N], ord;

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    cin >> n;
    for (ll i = 1; i <= n; i++) {
        ll c;
        cin >> c;
        dp[i] = c ? 1 : -1;
    }
    for (ll i = 1; i < n; i++) {
        ll u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    ord.push_back(1);
    for (ll i = 0; i < (ll)ord.size(); i++) {
        ll u = ord[i];
        for (ll v : g[u])
            if (v != par[u]) {
                par[v] = u;
                ord.push_back(v);
            }
    }
    for (ll i = n - 1; i >= 1; i--) {
        ll v = ord[i];
        dp[par[v]] += max(0LL, dp[v]);
    }
    ans[1] = dp[1];
    for (ll i = 1; i < n; i++) {
        ll v = ord[i];
        ans[v] = dp[v] + max(0LL, ans[par[v]] - max(0LL, dp[v]));
    }
    for (ll u = 1; u <= n; u++) cout << ans[u] << ' ';
    cout << '\n';
}

ChạyInput5
1 0 1 1 0
1 2
2 3
2 4
4 5

Output2 2 2 2 1

$O(n)$.

Tiền tố, hậu tố

Khi phép gộp không có phép ngược (tích theo modulo không nguyên tố, max, GCD), không thể “bỏ đóng góp của một con” bằng phép chia hay phép trừ. Thay vào đó, ở mỗi đỉnh $u$ với các con $c_1, \dots, c_k$, tính tích tiền tố $pre_i$ của $i$ con đầu và tích hậu tố $suf_i$ từ con $i$ về cuối.

Phần “mọi con trừ $c_i$” là $pre_{i-1} \cdot suf_{i+1}$. Tổng kích thước các mảng này là tổng số con, tức $n - 1$, nên vẫn $O(n)$.

Ví dụCụm Đen Chứa Đỉnh

Đề. Cây $n \le 10^5$ đỉnh, số $m \le 10^9$ không nhất thiết nguyên tố. Với mỗi đỉnh $u$, đếm số tập đỉnh liên thông chứa $u$, theo modulo $m$.

Lời giải. Từ lá lên: $dp_u$, tích các $(dp_c + 1)$ trên mọi con $c$, là số tập liên thông chứa $u$ trong cây con của nó, vì mỗi con hoặc không lấy gì, hoặc lấy một tập liên thông chứa nó. Từ gốc xuống: $up_c$ là số tập liên thông chứa $u$ nằm ngoài cây con của $c$, bằng $(up_u + 1)$ nhân với tích các $(dp_{c'} + 1)$ trên mọi con $c' \ne c$ của $u$, với $up$ của gốc là 0. Đáp án của $c$ là $dp_c \cdot (up_c + 1)$: phần dưới chứa $c$, nhân với lựa chọn có lấy phía trên hay không.

Khác ví dụ trước, phép gộp là tích theo modulo $m$ không nhất thiết nguyên tố, nên không “chia ngược” được để bỏ một con. Thay vào đó dùng tích tiền tố và hậu tố của các con.

Bước 1: danh sách con. ch[u] dựng trong BFS để dùng lại cho tiền tố và hậu tố.

for (ll v : g[u])
    if (v != par[u]) {
        par[v] = u;
        ch[u].push_back(v);
        ord.push_back(v);
    }

Bước 2: từ lá lên.

for (ll i = n - 1; i >= 0; i--) {
    ll u = ord[i];
    dp[u] = 1;
    for (ll c : ch[u]) dp[u] = dp[u] * (dp[c] + 1) % m;
}

Bước 3: từ gốc xuống bằng tiền tố, hậu tố. Ở mỗi $u$, pre[j] là tích của $j$ con đầu, suf[j] là tích từ con $j$ tới cuối. Tích mọi con trừ con thứ $j$ là pre[j] * suf[j + 1], đúng với mọi mô-đun.

for (ll i = 0; i < n; i++) {
    ll u = ord[i], k = ch[u].size();
    vector<ll> pre(k + 1, 1), suf(k + 2, 1);
    for (ll j = 0; j < k; j++) pre[j + 1] = pre[j] * (dp[ch[u][j]] + 1) % m;
    for (ll j = k - 1; j >= 0; j--) suf[j] = suf[j + 1] * (dp[ch[u][j]] + 1) % m;
    for (ll j = 0; j < k; j++) up[ch[u][j]] = (up[u] + 1) % m * pre[j] % m * suf[j + 1] % m;
}

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

const ll N = 1e5 + 5;
ll n, m, par[N], dp[N], up[N];
vector<ll> g[N], ch[N], ord;

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    cin >> n >> m;
    for (ll i = 1; i < n; i++) {
        ll u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    ord.push_back(1);
    for (ll i = 0; i < (ll)ord.size(); i++) {
        ll u = ord[i];
        for (ll v : g[u])
            if (v != par[u]) {
                par[v] = u;
                ch[u].push_back(v);
                ord.push_back(v);
            }
    }
    for (ll i = n - 1; i >= 0; i--) {
        ll u = ord[i];
        dp[u] = 1;
        for (ll c : ch[u]) dp[u] = dp[u] * (dp[c] + 1) % m;
    }
    for (ll i = 0; i < n; i++) {
        ll u = ord[i], k = ch[u].size();
        vector<ll> pre(k + 1, 1), suf(k + 2, 1);
        for (ll j = 0; j < k; j++) pre[j + 1] = pre[j] * (dp[ch[u][j]] + 1) % m;
        for (ll j = k - 1; j >= 0; j--) suf[j] = suf[j + 1] * (dp[ch[u][j]] + 1) % m;
        for (ll j = 0; j < k; j++) up[ch[u][j]] = (up[u] + 1) % m * pre[j] % m * suf[j + 1] % m;
    }
    for (ll u = 1; u <= n; u++) cout << dp[u] * (up[u] + 1) % m << '\n';
}

ChạyInput5 100
1 2
2 3
2 4
4 5

Output7
12
7
10
6

Mỗi đỉnh tốn số con của nó: $O(n)$.

Code trên khai báo hai vector mới ở mỗi đỉnh cho gọn. Khi cây cỡ triệu đỉnh, việc cấp phát ấy đáng kể:

Dùng hai mảng toàn cục chung cho mọi đỉnh nhanh hơn khoảng 2 lần. Vòng lặp chạy cỡ triệu lần thì tránh khai báo vector bên trong.