洛谷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;
}