换根DP 牛客练习赛99

已知以某一个点为根节点的树的总深度之和，如何求以其它节点为根的总深度之和？

如图：x是y的根节点

f[x]以x为根节点的树的总深度之和，把它的一个子节点y拿来当根节点，f[y]=f[x]-cnt[y]+(n-cnt[y])。

D 礼物

 代码：

#include<bits/stdc++.h>
using namespace std;
const int N=1e6+10;
typedef long long ll;
ll f[N],s[N];
int h[N],ne[N*2],e[N*2],idx;
int n;
void add(int a,int b)
{
    e[idx]=b,ne[idx]=h[a],h[a]=idx++;
}
void dfs1(int u,int fa)
{
    s[u]=1;
    for(int i=h[u];i!=-1;i=ne[i])
    {
        int j=e[i];
        if(j==fa) continue;
        dfs1(j,u);
        s[u]+=s[j];
    }
    for(int i=h[u];i!=-1;i=ne[i])
    {
        int j=e[i];
        if(j==fa) continue;
        f[1]+=(ll)u*s[j]*(s[u]-s[j]);//以u为根的子树中，以j为根的这部分子树里的所有点与剩下的的所有点的lca都是u
    }                                //所以有s[j]*(s[u]-s[j])种搭配
    f[1]+=(ll)u*s[u];//以u为根的所有子树的节点与u的lca
}
//换根DP，
void dfs2(int u,int fa)
{
    for(int i=h[u];i!=-1;i=ne[i])
    {
        int j=e[i];
        if(j==fa) continue;
        int tmp=s[j];
        s[u]=n-s[j];
        s[j]=n;
        f[j]=f[u]+(ll)j*(n-tmp)*tmp*2-(ll)u*tmp*(n-tmp)*2;
        //j*(n-tmp)*tmp：这一部分是深度加1的
        //u*tmp*(n-tmp)：这一部分是深度减1的
        dfs2(j,u);
        s[u]=n;
        s[j]=tmp;
    }
}
int main()
{
    memset(h,-1,sizeof(h));
    scanf("%d",&n);
    int u,v;
    for(int i=1;i<n;i++)
    {
        scanf("%d%d",&u,&v);
        add(u,v),add(v,u);
    }
    dfs1(1,0);
    dfs2(1,0);
    int idx=1;
    for(int i=1;i<=n;i++)
        if(f[i]>f[idx]) idx=i;
    printf("%d %lld\n",idx,f[idx]);
    return 0;
}