洛谷题单指南-图论之树-P5666 [CSP-S2019] 树的重心

原题链接：https://www.luogu.com.cn/problem/P5666

题意解读：计算树中所有边分割成两个子树各自的重心编号之和的和。

解题思路：

先分析一下重心的性质：

1、重心一定存在，有一个或者两个

2、找重心时，一定沿着根节点往重链方向找，如果重儿子所在子树大小>n/2，则一定要往重儿子方向找，一直到最后一个重儿子，使得重儿子以上的节点数<=n/2，这个重儿子即为重心

证明：首先，从根节点开始往下找，如果重儿子所在子树大小超过n/2，重心如果不在重儿子所在子树，必然导致重心分割的子树大小有大于n/2的情况，产生矛盾；其次，如果重心不是最后一个符合条件的重儿子，儿子后面还有重儿子可以让其以上的节点数<=n/2，那么重心以下的节点数必然>n/2，不符合重心的定义；最后，重儿子子树大小不超过n/2，重儿子以上节点数也不超过n/2，符合重心定义。

3、如果存在两个重心，另一个重心一定是上面找到的第一个重心的父节点，判断条件为n为偶数，且第一重儿子子树大小正好是n/2。

证明：重心既然在重链上，那么第二个重心只是第一个重心的子节点或者父节点，只需要证明不能是子节点。如果是子节点，说明第一个重心不是最后一个符合条件的重儿子，产生矛盾，因此必须是父节点。两个相邻的重心都可以将节点分为两部分，节点数为n，设分隔方式为a-1-1-b，两个1表示重心，有a+1<=n/2, b+1<=n/2，a+b+2<=n，又因为a+b+2就是所有节点数n，因此只能a=b且n为偶数才能满足，也就是两个重心分出的两个子树大小相同。

有了以上分析，找重心可以从根节点开始，沿着重链找最后一个使得上面节点数<=n/2的重儿子。

要找每条边分割之后两个子树的重心，就要在DFS每条边u->v的过程中，一方面沿着v往下找，另一方面沿着u往上找。

需要两次dfs，第一次dfs预处理出节点子树大小、父节点、重儿子等基本信息；第二次dfs枚举每条边，每条边的父子节点为根将树划分为两个子树，对子节点沿着重儿子找重心，对父节点要进行换根操作，之后再沿着父节点的重儿子找重心。

需要注意，如果父节点的重儿子就是子节点，不能走子节点的路，要继续从次重儿子找父节点所在子树的重心。

另外由于找重心的过程需要沿着重链枚举重儿子，如果极端情况复杂度是O(n)，n-1条边都处理一次，总的复杂度是O(n^2)，需要优化沿着重链找重儿子的过程，这里可以通过倍增来实现，因为一定能找到最后一个满足条件的重儿子，再往下找重儿子以上节点数就会超过n/2。这样总体复杂度为O(nlogn)。

100分代码：

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

const int N = 300005;

vector<int> g[N];
int siz[N]; // 子树大小
int fa[N]; // 父节点
int son[N][20]; // son[u][i]表示u的第2^i个重儿子
int t, n;
long long ans;

void init(int u)
{
    for(int i = 1; i <= 19; i++)
    {
        son[u][i] = son[son[u][i-1]][i-1];
    }
}

void dfs1(int u, int p)
{
    siz[u] = 1;
    fa[u] = p;
    for(int v : g[u])
    {
        if(v == p) continue;
        dfs1(v, u);
        siz[u] += siz[v];
        if(siz[v] > siz[son[u][0]]) son[u][0] = v;
    }
    init(u);
}

void find_centroid(int u)
{
    int cnt = siz[u];
    for(int i = 19; i >= 0; i--) // 倍增找重心
    {
        if(son[u][i] != 0 && cnt - siz[son[u][i]] <= cnt / 2)
        {
            u = son[u][i];
        }
    }
    ans += u;
    if(cnt % 2 == 0 && siz[u] == cnt / 2) 
    {
        ans += fa[u]; // 有两个重心的情况
    }
}

void dfs2(int u, int p)
{
    int sizu = siz[u];
    int sonu = son[u][0];
    int s1 = 0, s2 = 0; // s1表示u的重儿子，s2表示u的次重儿子
    for(int v : g[u])
    {
        //注意这里要考虑向上的节点，u要往各个方向找重儿子和次重儿子
        //因为向上的节点已经换过根，节点大小已经更新过了
        //因此不能加这句：if(v == p) continue;
        if(siz[v] >= siz[s1]) s2 = s1, s1 = v; //这里的大小关系一定不能搞错
        else if(siz[v] > siz[s2]) s2 = v; //这里的大小关系一定不能搞错
    }
    for(int v : g[u])
    {
        if(v == p) continue;
        find_centroid(v); // 找以v为根的子树重心

        //换根
        siz[u] = n - siz[v]; //u为根时的子树大小
        fa[u] = v; //u的父节点变成v
        if(s1 == v) son[u][0] = s2;
        else son[u][0] = s1;
        init(u);
        find_centroid(u); // 找以u为根的子树重心
        dfs2(v, u);
        // 回溯
        son[u][0] = sonu;
        init(u);
        siz[u] = sizu;
        fa[u] = p;
    }
}

int main() 
{
    cin >> t;
    while(t--)
    {
        cin >> n;
        memset(g, 0, sizeof(g));
        memset(son, 0, sizeof(son));
        memset(siz, 0, sizeof(siz));
        memset(fa, 0, sizeof(fa));
        for(int i = 1; i < n; i++)
        {
            int u, v;
            cin >> u >> v;
            g[u].push_back(v);
            g[v].push_back(u);
        }
        ans = 0; 
        dfs1(1, 0);
        dfs2(1, 0);
        cout << ans << endl;
    }

    return 0;
}