CF1187E Tree Painting

Tree Painting - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

Solution

容易得出结论：可选择的只有第一个染色点，且不妨将该点视作根结点。该点确定后，本题答案也就确定了为所有子树的大小之和。现在要求的是如何选择根结点，使得答案最大。

设 $f_i$ 为以 $i$ 为根的子树对答案的贡献。当 $i$ 为树根时，答案为 $f_i$。$siz_i$ 为 $1$ 为树根时，以 $i$ 为根的子树的大小。不妨一开始以 $1$ 为第一个染色点，然后进行换根 dp，按 dfs 序选择新的根。设 $i$ 为 $1$ 的儿子，则可知

$f_i=n+\sum \limits_{k\in son(1) \backslash i} f_k+n-siz_i+\sum \limits_{k\in son(i)} f_k$

又因为

$\sum \limits_{k\in son(i)} f_k=f_i-siz_i$

和

$\sum \limits_{k\in son(1)} f_k=f_1-n$

联立三式得到

$f_i=n+f_1-2siz_i$

这个例子是把根从 $1$ 换到子结点 $i$。同理，所有的 $f_i$ 均可按照同样的方法由父结点的 $f$ 得到。最大的 $f_i$ 为所求。

Code

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

const ll N = 2e5 + 10;

ll n, f[N], ans, siz[N];
vector<ll> g[N];

void dfs(ll u, ll fa)
{
    siz[u] = 1;
    for (auto v : g[u]) {
        if (v == fa) continue;
        dfs(v, u);
        siz[u] += siz[v];
        f[u] += f[v];
    }
    f[u] += siz[u];
}

void dp(ll u, ll fa)
{
    ans = max(ans, f[u]);
    for (int v : g[u]) {
        if (v == fa) continue;
        f[v] = n + f[u] - 2 * siz[v];
        dp(v, u);
    }
}

int main()
{
    scanf("%lld", &n);
    for (int i = 1; i < n; i++) {
        ll u, v;
        scanf("%lld%lld", &u, &v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1, 0); ans = f[1]; dp(1, 0);
    printf("%lld", ans);
    return 0;
}