[BZOJ 3631][JLOI2014]松鼠的新家 dfs

维护懒惰标记就好了.

#include<cstdio>
#include<cstring>
#include<iostream>
#define maxn 300021
using namespace std;
int n,a[maxn],head[maxn],tot=1,val[maxn];
int f[maxn][21],h[maxn],lz1[maxn],ans[maxn],lz2[maxn];
struct edge{int v,next;}e[maxn*2];
void adde(int a,int b){e[tot].v=b,e[tot].next=head[a];head[a]=tot++;}
void dfs(int u){
	for(int i=1;i<=20;i++)f[u][i]=f[f[u][i-1]][i-1];
	for(int v,i=head[u];i;i=e[i].next){
		if((v=e[i].v)==f[u][0])continue;
		f[v][0]=u;
		h[v]=h[u]+1;
		dfs(v);
	}
}
int lca(int a,int b){
	if(h[a]>h[b])swap(b,a);
	for(int i=20;i>=0;i--)if(h[f[b][i]]>=h[a])b=f[b][i];
	if(a==b)return a;
	for(int i=20;i>=0;i--){
		if(f[a][i]==f[b][i])continue;
		a=f[a][i],b=f[b][i];
	}return f[a][0];
}
void get(int u,int fa){
	for(int i=head[u],v;i;i=e[i].next){
		if((v=e[i].v)==fa)continue;
		get(v,u);
	}
	val[u]+=lz1[u];
	ans[u]+=val[u];
	val[u]+=lz2[u];
	val[fa]+=val[u];
}
int main(){
	scanf("%d",&n);
	for(int i=1;i<=n;i++)scanf("%d",a+i);
	int x,y,g,last;
	for(int i=1;i<n;i++){
		scanf("%d%d",&x,&y);
		adde(x,y),adde(y,x);
	}h[1]=1;
	dfs(1);
	last=a[1];
	for(int i=1;i<=n;i++){
		ans[i]--;
		x=last,y=a[i];g=lca(x,y);
		if(h[x]>h[y])swap(x,y);
		if(x==g)lz1[y]++,lz2[x]--;
		else lz1[x]++,lz1[y]++,lz2[g]-=2,ans[g]--;
		last=last^x^y;
	}get(1,0);
	for(int i=1;i<=n;i++)printf("%d\n",ans[i]);
	return 0;
}