洛谷P5666 [CSP-S2019] 树的重心题解

题目链接：P5666 [CSP-S2019] 树的重心

题意：

小简单正在学习离散数学，今天的内容是图论基础，在课上他做了如下两条笔记：

一个大小为 $n$ 的树由 $n$ 个结点与 $n - 1$ 条无向边构成，且满足任意两个结点间有且仅有一条简单路径。在树中删去一个结点及与它关联的边，树将分裂为若干个子树；而在树中删去一条边（保留关联结点，下同），树将分裂为恰好两个子树。

对于一个大小为 $n$ 的树与任意一个树中结点 $c$，称 $c$ 是该树的重心当且仅当在树中删去 $c$ 及与它关联的边后，分裂出的所有子树的大小均不超过 $\lfloor \frac{n}{2} \rfloor$（其中 $\lfloor x \rfloor$ 是下取整函数）。对于包含至少一个结点的树，它的重心只可能有 1 或 2 个。

课后老师给出了一个大小为 $n$ 的树 $S$，树中结点从 $1 \sim n$ 编号。小简单的课后作业是求出 $S$ 单独删去每条边后，分裂出的两个子树的重心编号和之和。即：
$\sum_{(u,v) \in E} \left( \sum_{1 \leq x \leq n \atop 且 x 号点是 S'_u 的重心} x + \sum_{1 \leq y \leq n \atop 且 y 号点是 S'_v 的重心} y \right)$
上式中，$E$ 表示树 $S$ 的边集，$(u,v)$ 表示一条连接 $u$ 号点和 $v$ 号点的边。$S’_u$ 与 $S’_v$ 分别表示树 $S$ 删去边 $(u,v)$ 后，$u$ 号点与 $v$ 号点所在的被分裂出的子树。

小简单觉得作业并不简单，只好向你求助，请你教教他。

输入格式：

本题包含多组测试数据。

第一行一个整数 $T$ 表示数据组数。

接下来依次给出每组输入数据，对于每组数据：

第一行一个整数 $n$ 表示树 $S$ 的大小。

接下来 $n − 1$ 行，每行两个以空格分隔的整数 $u_i$，$v_i$，表示树中的一条边 $(u_i,v_i)$。

输出格式：

共 $T$ 行，每行一个整数，第 $i$ 行的整数表示：第 $i$ 组数据给出的树单独删去每条边后，分裂出的两个子树的重心编号和之和。

数据范围：

$1 \leq T \leq 5 , 1 \leq u_i,v_i \leq n$。保证给出的图是一个树。

套路地考察每个点作为重心的次数。考虑将树的重心提作根 $\rm root$ （有两个就随便选一个）。

对于任意点 $u~(u \ne \mathrm{root})$ ，若 $u$ 是割掉某条边后的重心，则这条边一定不在 $u$ 的子树内。

记割掉某条边 $e$ 后，另一个树的大小为 $S$ ，并记

$g_u = \max_{v \in \mathrm{son}(u)}\left\{ \mathrm{size}(v)\right\}$

那么根据 $u$ 为重心可知

$\begin{cases} 2(n-S-\mathrm{size}(u)) \le n - S \\[8pt] 2g_u \le n -S \end{cases}$

所以

$n-2\ \mathrm{size}(u) \leq S \leq n-2 g_u$

即我们需要找到对于点 $u$ 有多少条边 $e$ 满足

$n-2\ \mathrm{size}(u) \leq S \leq n-2 g_u$
边不在 $u$ 的子树内。

仅考虑第一个条件的话，考虑 dfs 并用树状数组维护切掉边 $e = (u, fa(u))$ 后，

每一个 $S$ 的值有多少个，这样对于每个 $u$ 我们只需要查询一下区间和即可。

考虑去掉不满足第二个条件的 $e$ ，考虑再拿一个树状数组按 dfs 序动态维护经过每个点的 $S$ 的值有多少个

这样对于每个点，在其子树内的 $e$ 就是进入时减去回溯时的区间和的差。

那么怎么统计 $u=\rm root$ 的贡献呢？设 $u$ 的儿子中 $\rm size$ 最大的节点是 $x$ ，次大的是 $y$ 。

若 $e$ 在 $x$ 的子树中，则有

$2\ \mathrm{size}(y) \le n -S$

即

$S \le n - 2\ \mathrm{size}(y)$

否则有

$2\ \mathrm{size}(x) \le n - S$

即

$S \le n - 2\ \mathrm{size}(x)$

直接维护即可。

时间复杂度 $\mathcal{O}(n \log n)$

代码：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
void up(int &x, int y) { x < y ? x = y : 0; }
void down(int &x, int y) { x > y ? x = y : 0; }
#define rep(i, a, b) for(int i = (a), i##END = (b); i <= i##END; i++)
#define Rep(i, a, b) for(int i = (a), i##END = (b); i >= i##END; i--)
#define pb push_back
#define lowbit(x) ((x) & -(x))
#define N ((int)(3e5 + 15))

vector<int> vec[N];
int n, rt, res, A, B, sz[N], g[N], tr1[N], tr2[N]; char vip[N];
void add(int *tr, int x, int v) { for(int i = ++x; i <= n + 1; i += lowbit(i)) tr[i] += v; }
int qry(int *tr, int x, int r = 0) { for(int i = ++x; i; i -= lowbit(i)) { r += tr[i]; } return r; }
void dfs1(int u, int fa)
{
    sz[u] = 1, g[u] = 0; bool ok = true;
    for(int v : vec[u]) if(v != fa)
    {
        dfs1(v, u); sz[u] += sz[v]; up(g[u], sz[v]);
        if(sz[v] > n / 2) ok = false;
    }
    if(n - sz[u] > n / 2) ok = false;
    if(ok) rt = u;
}
void dfs2(int u, int fa)
{
    add(tr1, sz[fa], -1); add(tr1, n - sz[u], 1);
    if(u != rt)
    {
        res += u * (qry(tr1, n - 2 * g[u]) - qry(tr1, n - 2 * sz[u] - 1));
        res += u * (qry(tr2, n - 2 * g[u]) - qry(tr2, n - 2 * sz[u] - 1));
        if(!vip[u] && vip[fa]) vip[u] = true;
        res += rt * (sz[u] <= n - 2 * sz[vip[u] ? B : A]);
    }
    add(tr2, sz[u], 1);
    for(int v : vec[u]) if(v != fa) dfs2(v, u);
    add(tr1, sz[fa], 1); add(tr1, n - sz[u], -1);
    if(u != rt) { res -= u * (qry(tr2, n - 2 * g[u]) - qry(tr2, n - 2 * sz[u] - 1)); }
}
void clear()
{
    res = A = B = rt = 0; 
    rep(i, 0, n + 5) { vec[i].clear(); tr1[i] = tr2[i] = vip[i] = 0; }
}
void solve()
{
    cin >> n; clear();
    rep(i, 1, n - 1) { static int u, v; cin >> u >> v; vec[u].pb(v); vec[v].pb(u); }
    dfs1(1, 0); dfs1(rt, 0);
    for(int v : vec[rt]) { if(sz[v] > sz[B]) B = v; if(sz[B] > sz[A]) swap(A, B); }
    rep(i, 0, n) add(tr1, sz[i], 1);
    vip[A] = true; dfs2(rt, 0); cout << res << '\n';
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    int qwq; cin >> qwq; while(qwq--) solve();
    return 0;
}