树形dp之换根dp

#include<bits/stdc++.h>
using namespace std;
#define N 1000005
#define int long long
int n,cnt=0,ans=0,h[N]= {0};//cnt为边的总数，h存各点头节点位置 ，ans为我们要求的节点
int si[N]= {0},d[N]= {0},f[N]= {0};//si节点子节点个数，这里为以1为根节点时，d为节点深度，f为题中所描述的函数
struct edge {
    int to,next;
} e[2*N];     //链式前向星存图
void  add(int x,int y) {
    e[++cnt].next=h[x];
    e[cnt].to=y;
    h[x]=cnt;
}//链式前向星加边
void dfs(int fa,int s) {
    si[s]=1;
    d[s]=d[fa]+1;
    for(int i=h[s]; i; i=e[i].next) {
        if(e[i].to==fa)continue;
        dfs(s,e[i].to);
        si[s]+=si[e[i].to];
    }
}
void dfs2(int fa,int s) {
    for(int i=h[s]; i; i=e[i].next) {
        if(e[i].to==fa)continue;
        f[e[i].to]+=f[s]+n-2*si[e[i].to];
        /*如果节点从该节点父节点移动至该节点，他的值为父节点的值加上该点父节点的父（包括爷爷）节点的 个数
        减去该节点儿子的个数（即si[e[i]].to）,父节点个数等于n减去他的儿子的值
        即：f[s]+n-si[e[i].to]-si[e[i].to]]*/
        dfs2(s,e[i].to);
    }
}
signed main() {
    cin>>n;
    for(int i=1; i<n; i++) {
        int x,y;
        cin>>x>>y;
        add(x,y);
        add(y,x);
    }
    dfs(0,1);//第一个dfs进行初始化，在1为根节点的情况下，各节点的子（包含孙子节点）节点个数 和深度
    for(int i=1; i<=n; i++)f[1]+=d[i];//这里把各深度相加即得到以1为根节点情况下的值
    dfs2(0,1);//此处为深搜进行的一个动态规划
    f[0]=0;
    for(int i=1; i<=n; i++)if(f[ans]<f[i])ans=i;
    cout<<ans;
    return 0;
}