CSP2019 Day2T3 树的重心

显然如果我们直接枚举断哪条边，剩下的两颗树的重心编号非常不好求，一个常见的想法是我们反过来，考虑每个节点有多少种情况作为树的重心，计算每个点对答案的贡献。

下面我们就需要大力分类讨论了。假设我们现在考虑计算贡献的节点为 (x)，令以其儿子为根的子树树的最大的子树大小为 (mx_x)，次大值为 (se_x)，以 (x) 为根的子树大小为 (S_x)，假设 (x) 的一个祖先 (f) 断开了其与父亲的边。那么如果 (x) 要作为以 (f) 为根的树的重心，那么需要满足下面一个条件：

[egin{cases} S_f - S_x le lfloor frac{S_f}{2} floor\ mx_x le lfloor frac{S_f}{2} floor end{cases}\]

因为我们需要统计的是满足条件的 (y) 的形式，因此我们尽量将原式化简成关于 (y) 的一个范围。实际上，上面那个柿子能改写成 (2 imes S_f - 2 imes S_x le S_f)，移项可得 (S_f le 2 imes S_x)，类似地下面那个柿子有 (S_f ge 2 imes mx_x)，于是我们只需要统计一个点 (x) 到根的路径上有多少个点满足 (2 imes mx_x le S_f le 2 imes S_x) 即可。这个我们可以使用树状数组来做，进入每个点是将 (S_x) 加入树状数组，出去时删除，那么递归到每个点时的树状数组就是加入 (x) 到根路径上点的树状数组了，如果硬核数据结构也可以用主席树实现。

再考虑第二种情况，断掉 (x) 子树内的一个点 (y) 与其父亲的连边的情况。又要分两种情况，当 (y) 在 (x) 的重儿子内时，需要满足：

[egin{cases} se_x le lfloor frac{S_1 - S_y}{2} floor\ mx_x - S_y le lfloor frac{S_1 - S_y}{2} floor\ S_1 - S_x le lfloor frac{S_1 - S_y}{2} floor end{cases}]

同理上面的化简可得：(2 imes mx_x - S_1 le S_y le min{S_1 - 2 imes se_x, 2 imes S_x - S_1})，实际上这里我们需要的是统计子树内权值在某个范围的数量，可以使用线段树合并解决。

再考虑如果 (y) 在 (x) 的非重儿子内部，那么我们需要满足：

[egin{cases} mx_x le lfloor frac{S_1 - S_y}{2} floor\ S_1 - S_x le lfloor frac{S_1 - S_y}{2} floor end{cases}]

可以得出 (S_y le min{S_1 - 2 imes mx_x, 2 imes S_x - S_1})，同样是查询子树信息，用线段树合并解决。

我们再考虑最后一种情况，当不为 (x) 的祖先且不在 (x) 子树内的一个点 (y) 断掉了其与父亲的连边时，需要满足：

[egin{cases} S_1 - S_y - S_x le lfloor frac{S_1 - S_y}{2} floor\ mx_x le lfloor frac{S_1 - S_y}{2} floor end{cases}]

化简可得：(S_1 - 2 imes S_x le S_y le S_1 - 2 imes mx_x)。

这个时候我被难住了，我们怎么知道不包含 (x) 的祖先和其子树内的点满足条件的所有 (y) 呢？其实很简单，上面我们已经知道了怎么统计在 (x) 的子树和其子树内满足条件的方法了，我们直接考虑容斥即可。考虑使用全局满足条件的点数减去在 (x) 的祖先和其子树内满足条件的点数即可，这个全局满足条件的点可以使用 (ton) 求出。

代码细节比较多，想清楚再码。

#include<bits/stdc++.h>
using namespace std;
#define N 300000 + 5
#define M 6000000 + 5
#define ls t[p].l
#define rs t[p].r
#define mid (l + r) / 2
#define lowbit(x) (x & (-x))
#define rep(i, l, r) for(int i = l; i <= r; ++i)
#define Next(i, u) for(int i = h[u]; i; i = e[i].next)
struct edge{
    int v, next;
}e[N << 1];
struct tree{
    int l, r, sum;
}t[M];
long long ans;
int T, n, u, v, tot, cnt, h[N], c[N], s[N], rt[N], mx[N], se[N], ton[N];
int read(){
    char c; int x = 0, f = 1;
    c = getchar();
    while(c > '9' || c < '0'){ if(c == '-') f = -1; c = getchar();}
    while(c >= '0' && c <= '9') x = x * 10 + c - '0', c = getchar();
    return x * f;
}
void Add(int u, int v){
    e[++cnt].v = v, e[cnt].next = h[u], h[u] = cnt;
    e[++cnt].v = u, e[cnt].next = h[v], h[v] = cnt;
}
void add(int p, int k){
    for(int i = p; i <= n; i += lowbit(i)) c[i] += k;
}
int ask(int p){
    int ans = 0; p = min(p, n);
    for(int i = p; i >= 1; i -= lowbit(i)) ans += c[i];
    return ans;
}
void update(int &p, int l, int r, int x, int y, int k){
    if(!p) p = ++tot; t[p].sum += k;
    if(l == r) return;
    if(mid >= x) update(ls, l, mid, x, y, k);
    if(mid < y) update(rs, mid + 1, r, x, y, k);
}
int query(int p, int l, int r, int x, int y){
    if(!p || l > r) return 0;
    if(l >= x && r <= y) return t[p].sum;
    int ans = 0;
    if(mid >= x) ans += query(ls, l, mid, x, y);
    if(mid < y) ans += query(rs, mid + 1, r, x, y);
    return ans;
}
void Merge(int &p, int k, int l, int r){
    if(!p || !k){ p = p + k; return;}
    if(l == r){ t[p].sum += t[k].sum; return;}
    Merge(ls, t[k].l, l, mid), Merge(rs, t[k].r, mid + 1, r);
    t[p].sum = t[ls].sum + t[rs].sum;
}
void dfs1(int u, int fa){
    s[u] = 1;
    Next(i, u){
        int v = e[i].v; if(v == fa) continue;
        dfs1(v, u), s[u] += s[v];
        if(s[v] >= mx[u]) se[u] = mx[u], mx[u] = s[v];
        else if(s[v] > se[u]) se[u] = s[v];
    }
}
void dfs2(int u, int fa){
    int tmp = ask(2 * s[u]) - ask(2 * mx[u] - 1) - (ask(s[1] - 2 * mx[u]) - ask(s[1] - 2 * s[u] - 1));
    if(mx[u] <= s[u] / 2 && u != 1) ++tmp;
    if(s[1] >= 2 * mx[u] && s[1] <= 2 * s[u] && u != 1) --tmp;
    add(s[u], 1);
    Next(i, u){
        int v = e[i].v; if(v == fa) continue;
        dfs2(v, u), Merge(rt[u], rt[v], 1, n);
        if(s[v] == mx[u]) tmp += query(rt[v], 1, n, max(1, 2 * mx[u] - s[1]), min(s[1] - 2 * se[u], 2 * s[u] - s[1]));
        else tmp += query(rt[v], 1, n, 1, min(s[1] - 2 * mx[u], 2 * s[u] - s[1]));
    }
    add(s[u], -1), update(rt[u], 1, n, s[u], s[u], 1);
    tmp -= query(rt[u], 1, n, max(1, s[1] - 2 * s[u]), s[1] - 2 * mx[u]);
    if(s[1] - 2 * mx[u] >= 1) tmp += ton[s[1] - 2 * mx[u]];
    if(s[1] - 2 * s[u] >= 1) tmp -= ton[s[1] - 2 * s[u] - 1];
    ans += 1ll * tmp * u;
}
int main(){
    T = read();
    while(T--){
        memset(h, 0, sizeof(h)), memset(rt, 0, sizeof(rt));
        memset(s, 0, sizeof(s)), memset(mx, 0, sizeof(mx));
        memset(se, 0, sizeof(se)), memset(ton, 0, sizeof(ton));
        rep(i, 1, tot) t[i].l = t[i].r = t[i].sum = 0;
        tot = cnt = ans = 0;
        n = read();
        rep(i, 1, n - 1) u = read(), v = read(), Add(u, v);
        dfs1(1, 0);
        rep(i, 1, n) ++ton[s[i]];
        rep(i, 1, n) ton[i] = ton[i - 1] + ton[i];
        dfs2(1, 0);
        printf("%lld
", ans);
    }
    return 0;
}