P5666 [CSP-S 2019] 树的重心

P5666 [CSP-S 2019] 树的重心

大意

删去树中每条边，树分裂为两个子树，求所有边对应的两子树重心值之和的总和。

思路

首先考虑 \(\mathcal{O}(n ^ 2)\) 的解法。

枚举每条边，然后分别统计左右子树的重心，然后想加即可。

然后考虑优化，显然我们的复杂度浪费在了找重心的地方。

“一棵有根树的重心一定在根结点所在的重链上。一棵树的重心一定是该树根结点重子结点对应子树的重心的祖先。”——OIwiki

有关树的重心的内容见 树的重心

重心必在节点的「重儿子链」上（重儿子指子树大小最大的子节点），因此可通过倍增沿重儿子链快速定位重心候选。

删除边 \((u, v)\) 后，\(u\) 所在的大子树（大小 \(n - sz_v\)）需要「换根」更新重儿子（原重儿子可能是 \(v\)，需替换为次重儿子或父方向），而 \(v\) 所在的小子树（大小 \(sz_v\))）是原树的子树，无需换根。

节点最大子树（含父方向）≤ 总大小 / \(2\) 则为重心。

递归算每个节点的子树大小，记录重 / 次重儿子（子树最大 / 次大的子节点），构建倍增数组（沿重儿子链跳，\(\mathcal{O}(\log n)\) 找重心）。

然后累加答案即可。

代码

#include<iostream>
#include<cstring>
using namespace std;

const int MAXN = 3e5 + 5;
const int LOG = 18;

struct node{
    int to, nxt;
}e[MAXN * 2];

int h[MAXN], tot;
int T, n;
int son1[MAXN], son2[MAXN];
int f[MAXN], up[MAXN][LOG + 1];
int son3[MAXN];
int fu[MAXN];
int sz[MAXN], csz[MAXN];
long long ans;

void add(int x, int y){
	e[++ tot] = {y, h[x]};
    h[x] = tot;
}

void dfs1(int u, int fa){
    sz[u] = 1;
    f[u] = fa;
    son1[u] = son2[u] = 0;
    for(int i = h[u];i;i = e[i].nxt){
        int v = e[i].to;
        if(v == fa) continue;
        dfs1(v, u);
        sz[u] += sz[v];
        if(sz[v] > sz[son1[u]]){
            son2[u] = son1[u];
            son1[u] = v;
        }
        else if(sz[v] > sz[son2[u]]){
            son2[u] = v;
        }
    }
    up[u][0] = son1[u];
    for(int i = 1;i <= LOG;i ++){
        up[u][i] = up[up[u][i - 1]][i - 1];
    }
}

int check(int x, int sum, int maxx){
    if(!x) return 0;
    return (max(maxx, sum - csz[x]) <= sum / 2) ? x : 0;
}

void dfs2(int u, int fa){
    for(int i = h[u];i;i = e[i].nxt){
        int v = e[i].to;
        if(v == fa) continue;
        //处理大子树u
        csz[u] = n - sz[v];
        son3[u] = (son1[u] == v) ? son2[u] : son1[u];
        if(fa && csz[fa] > csz[son3[u]]) son3[u] = fa;
        up[u][0] = son3[u];
        for(int j = 1;j <= LOG;j ++){
            up[u][j] = up[up[u][j - 1]][j - 1];
        }
        int hav1 = u;
        for(int j = LOG;j >= 0;j --){
            if(csz[u] - csz[up[hav1][j]] <= csz[u] / 2){
                hav1 = up[hav1][j];
            }
        }
        ans += check(son3[hav1], csz[u], csz[son3[son3[hav1]]]);
        ans += check(hav1, csz[u], csz[son3[hav1]]);
        ans += check(fu[hav1], csz[u], csz[son3[fu[hav1]]]);
        //处理小子树v
        csz[v] = sz[v];
        int hav2 = v;
        for(int j = LOG;j >= 0;j --){
            if(csz[v] - csz[up[hav2][j]] <= csz[v] / 2){
                hav2 = up[hav2][j];
            }
        }
        ans += check(son1[hav2], csz[v], csz[son1[son1[hav2]]]);
        ans += check(hav2, csz[v], csz[son1[hav2]]);
        ans += check(fu[hav2], csz[v], csz[son1[fu[hav2]]]);
        //递归处理子节点，维护临时父方向
        fu[u] = v;
        dfs2(v, u);
        //恢复原状态
        son3[u] = son1[u];
        up[u][0] = son1[u];
        for(int j = 1;j <= LOG;j ++){
            up[u][j] = up[up[u][j - 1]][j - 1];
        }
        csz[u] = sz[u];
        fu[u] = f[u];
    }
    son3[u] = son1[u];
    up[u][0] = son1[u];
    for(int j = 1;j <= LOG;j ++){
        up[u][j] = up[up[u][j - 1]][j - 1];
    }
    csz[u] = sz[u];
    fu[u] = f[u];
}

void init(){
    tot = 0; ans = 0;
    memset(h, 0, sizeof(h));
    memset(son1, 0, sizeof(son1));
    memset(son2, 0, sizeof(son2));
    memset(son3, 0, sizeof(son3));
    memset(f, 0, sizeof(f));
    memset(fu, 0, sizeof(fu));
    memset(sz, 0, sizeof(sz));
    memset(csz, 0, sizeof(csz));
    memset(up, 0, sizeof(up));
}

int main(){
	ios::sync_with_stdio(0);
	cin.tie(0);
	cin >> T;
    while(T --){
        init();
		cin >> n;
        for(int i = 1;i < n;i ++){
            int u, v; cin >> u >> v;
            add(u, v), add(v, u);
        }
        dfs1(1, 0);
        memcpy(csz, sz, sizeof(sz));
        memcpy(son3, son1, sizeof(son1));
        memcpy(fu, f, sizeof(f));
        dfs2(1, 0);
        cout << ans << '\n';
    }
    return 0;
}