D34【模板】长链剖分 CF1009F Dominant Indices

CF1009F Dominant Indices - 洛谷

长链剖分：按子树深度把树拆分成互不相交的若干条链。

每个节点只有 1 个（或 0 个）重儿子，其余都是轻儿子。

每条链的链头是轻儿子，其余为重儿子。

// 长链剖分+树形DP O(n)
#include<bits/stdc++.h>
using namespace std;

const int N=1000010;
vector<int> e[N];
int n,son[N],len[N];
int buf[N],*p=buf,*f[N],ans[N];
//f[x][i]:在x子树内与x的距离为i的节点数

void dfs(int x,int fa){
  for(int y:e[x]){
    if(y==fa) continue;
    dfs(y,x);
    if(len[son[x]]<len[y]) son[x]=y;
  }
  len[x]=len[son[x]]+1;
}
void DP(int x,int fa){
  f[x][0]=1; //节点自己
  if(son[x]){
    f[son[x]]=f[x]+1; //共享内存
    DP(son[x],x);
    ans[x]=ans[son[x]]+1; //递推答案
  }
  for(int y:e[x]){
    if(y==fa||y==son[x]) continue;
    f[y]=p; p+=len[y]; //给y开头的链分配内存
    DP(y,x);
    for(int j=1; j<=len[y]; j++){
      f[x][j]+=f[y][j-1]; //累计x下面j层的节点数
      if(f[x][j]>f[x][ans[x]]) ans[x]=j; //更新答案
      if(f[x][j]==f[x][ans[x]]&&j<ans[x]) ans[x]=j; //更新答案
    }
  }
  if(f[x][ans[x]]==1) ans[x]=0; //修正答案
}
int main(){
  scanf("%d",&n);
  for(int i=1,x,y; i<n; i++){
    scanf("%d%d",&x,&y);
    e[x].push_back(y); e[y].push_back(x);
  }    
  dfs(1,0); //长链剖分
  f[1]=p; p+=len[1]; //给树根长链分配内存
  DP(1,0); //树形DP
  for(int i=1;i<=n;++i) printf("%d\n",ans[i]);
}