E84【模板】换根DP P3478 [POI2008] STA-Station

// 换根DP O(n)
#include <iostream>
#include <cstring>
#include <algorithm>
using namespace std;

const int N=1000005;
int n,ans;
int head[N],to[N<<1],ne[N<<1],idx;
void add(int x,int y){
  to[++idx]=y,ne[idx]=head[x],head[x]=idx;
}
long long sz[N],dep[N],f[N],s;

void dfs(int u,int fa){ //预处理
  dep[u]=dep[fa]+1;
  f[1]+=dep[u];
  sz[u]=1;
  for(int i=head[u];i;i=ne[i]){
    int v=to[i];
    if(v==fa) continue;
    dfs(v,u);
    sz[u]+=sz[v];
  }
}
void dfs2(int u,int fa){ //换根DP
  for(int i=head[u];i;i=ne[i]){
    int v=to[i];
    if(v==fa) continue;
    f[v]=f[u]-sz[v]+n-sz[v];
    dfs2(v,u);
  }
}
int main(){
  scanf("%d",&n);
  for(int i=1,x,y;i<n;i++){
    scanf("%d%d",&x,&y);
    add(x,y);add(y,x);
  }
  dfs(1,0);
  dfs2(1,0);
  for(int i=1;i<=n;i++)
    if(f[i]>s) s=f[i],ans=i;
  printf("%d\n",ans);
}

sz[x]: x子树的节点数；f[x]: 以x为根时所有结点的深度之和。

dfs预处理，求出每颗子树的大小sz[i]，以1号点为根时所有结点的深度之和f[1]。

dfs2换根时，f[x]→f[y]：在y的子树上的结点深度都减少了1，即-sz[y]；不在y的子树上的结点深度都增加了1，即n-sz[y]；所以f[y] = f[x] - sz[y] + n - sz[y] = f[x] + n - 2*sz[y]。

#include <bits/stdc++.h>
using namespace std;
const int N=1e6+10;
vector<int> G[N];
int n, sz[N], dep[N];
long long f[N];

void dfs(int x, int fa) 
{
    dep[x] = dep[fa] + 1;
    f[1] += dep[x];
    sz[x] = 1;
    for(int y : G[x]) if(y != fa) 
    {
        dfs(y, x);
        sz[x] += sz[y];
    }
}
void dfs2(int x, int fa) 
{
    for(int y : G[x]) if(y != fa) 
    {
        f[y] = f[x] + n - 2 * sz[y];
        dfs2(y, x);
    }

}
int main() 
{
    scanf("%d", &n);
    for(int i=1, x, y; i <= n-1; i++) 
    {
        scanf("%d%d", &x, &y);
        G[x].push_back(y);
        G[y].push_back(x);
    }

    dep[0] = 0; f[0] = 0;
    dfs(1, 0);
    dfs2(1, 0);

    long long ans = -1;
    int id;
    for(int i=1; i <= n; i++) if(f[i] > ans) ans = f[i], id = i;
    printf("%d\n", id);
    return 0;
}