D58 树的直径 树上前缀和 P2971 [USACO10HOL] Cow Politics G

// 树上前缀和 O(nlogn)
#include<bits/stdc++.h>
using namespace std;

#define N 200005
int h[N],to[N<<1],ne[N<<1],idx;
void add(int u,int v){to[++idx]=v,ne[idx]=h[u],h[u]=idx;}

int n,k,rt,a[N],p[N];
int fa[N][21],dep[N],mxd[N],point[N],ans[N];

void dfs(int x,int f){
  dep[x]=dep[f]+1; fa[x][0]=f;
  for(int i=1;i<=20;i++)fa[x][i]=fa[fa[x][i-1]][i-1];
  if(mxd[a[x]]<dep[x]) mxd[a[x]]=dep[x],point[a[x]]=x; //记录同种颜色最深的点  
  for(int i=h[x];i;i=ne[i]){
    int y=to[i];
    if(y!=f) dfs(y,x);
  }
}
int lca(int x,int y){
  if(dep[x]<dep[y]) swap(x,y); //让x更深
  for(int i=20;~i;i--)if(dep[fa[x][i]]>=dep[y]) x=fa[x][i]; //x向上跳到y的同一层
  if(x==y) return x;
  for(int i=20;~i;i--)if(fa[x][i]!=fa[y][i]) x=fa[x][i],y=fa[y][i]; //一起向上跳
  return fa[x][0];
}
signed main(){
  scanf("%d%d",&n,&k);
  for(int i=1;i<=n;++i){
    scanf("%d%d",&a[i],&p[i]); //颜色,父亲
    if(p[i]==0) rt=i;
    else add(i,p[i]),add(p[i],i);
  }
  
  dfs(rt,0); //倍增预处理 dep,fa,point 数组
  for(int x=1,y;x<=n;++x){
    y=point[a[x]];
    ans[a[x]]=max(ans[a[x]],dep[x]+dep[y]-2*dep[lca(x,y)]); //树上前缀和
  } 
  for(int i=1;i<=k;++i) printf("%d\n",ans[i]);
}

#include<bits/stdc++.h>//by scy
using namespace std;
const int N =2e5+10;
vector<int> G[N];
int n,k,rt;
int a[N],D,dep[N],st[N][25],maxdep[N],pos[N],ans[N];
//maxdep[x]表示组别为x的集合中离根节点最远的点的深度
//   pos[x]表示组别为x的集合中离根节点最远的点的编号
void dfs(int x,int xfa)
{
	dep[x]=dep[xfa]+1;
    if(maxdep[a[x]]<dep[x]) maxdep[a[x]]=dep[x],pos[a[x]]=x;
    st[x][0]=xfa;for(int i=1;i<=D;i++) st[x][i]=st[st[x][i-1]][i-1];
	for(int y:G[x])dfs(y,x);
}
int lca(int x,int y)
{
	if(dep[x]<dep[y]) swap(x,y);
	for(int i=D;i>=0;--i)if(dep[st[x][i]]>=dep[y]) x=st[x][i];
	if(x==y) return x;
	for(int i=D;i>=0;--i)if(st[x][i]!=st[y][i]) x=st[x][i],y=st[y][i];
	return st[x][0];
}
int dis(int x,int y){return dep[x]+dep[y]-2*dep[lca(x,y)];}
int main()
{
	scanf("%d%d",&n,&k);
	for(int i=1,x;i<=n;++i)
    {
		scanf("%d%d",&a[i],&x);
		if(x==0) rt=i;
		else G[x].push_back(i);
	}
	memset(maxdep,0,sizeof(maxdep));
    dep[0]=0;D=log2(n);
	dfs(rt,0);
    memset(ans,0,sizeof(ans));
    //同个组别的所有点最远距离的路径的其中一个端点必定是深度最大的点
	for(int i=1;i<=n;++i) ans[a[i]]=max(ans[a[i]],dis(pos[a[i]],i));
	for(int i=1;i<=k;++i) printf("%d\n",ans[i]);
 	return 0;
}