bzoj2243: [SDOI2011]染色 树链剖分

裸树剖。

#include<bits/stdc++.h>
using namespace std;
#define N 100010
#define M (l+r>>1)
#define P (k<<1)
#define S (k<<1|1)
#define L l,M,P
#define R M+1,r,S
#define Z int l=1,int r=n,int k=1
typedef int ds[N];
ds dp,num,p,size,son,tmp,top,w;
struct edge{
	edge* s;
	int v;
}e[N<<1],*back=e,*h[N];
void add(int u,int v){
	h[u]=&(*back++=(edge){h[u],v});
	h[v]=&(*back++=(edge){h[v],u});
}
void dfs1(int u){
	size[u]=1;
	dp[u]=dp[p[u]]+1;
	int s=0;
	for(edge* i=h[u];i;i=i->s)
		if(i->v!=p[u]){
			p[i->v]=u;
			dfs1(i->v);
			size[u]+=size[i->v];
			if(s<size[i->v])
				s=size[son[u]=i->v];
		}
}
void dfs2(int u){
	static int cnt;
	tmp[num[u]=++cnt]=w[u];
	if(size[u]!=1){
		top[son[u]]=top[u];
		dfs2(son[u]);
	}
	for(edge* i=h[u];i;i=i->s)
		if(i->v!=p[u]&&i->v!=son[u])
			dfs2(top[i->v]=i->v);
}
typedef int node[1<<18];
node a,d,u,v;
int n;
void apply(int s,int k){
	a[k]=u[k]=v[k]=s;
	d[k]=1;
}
void devolve(int k){
	if(~a[k]){
		apply(a[k],P);
		apply(a[k],S);
		a[k]=-1;
	}
}
void update(int k){
	u[k]=u[P];
	v[k]=v[S];
	d[k]=d[P]+d[S]-(v[P]==u[S]);
}
void build(Z){
	if(l==r)
		apply(tmp[l],k);
	else{
		a[k]=-1;
		build(L);
		build(R);
		update(k);
	}
}
void A(int d,int s,int t,Z){
	if(s==l&&t==r)
		apply(d,k);
	else{
		devolve(k);
		if(t<=M)
			A(d,s,t,L);
		else if(s>M)
			A(d,s,t,R);
		else{
			A(d,s,M,L);
			A(d,M+1,t,R);
		}
		update(k);
	}
}
int Q(int s,int t,Z){
	if(s==l&&t==r)
		return d[k];
	devolve(k);
	return t<=M?Q(s,t,L)
	:s>M?Q(s,t,R)
	:Q(s,M,L)+Q(M+1,t,R)
	-(v[P]==u[S]);
}
int U(int i,Z){
	if(l==r)
		return a[k];
	devolve(k);
	return i<=M?U(i,L):U(i,R);
}
void color(int d,int s,int t){
	while(top[s]!=top[t]){
		if(dp[top[s]]<dp[top[t]])
			swap(s,t);
		A(d,num[top[s]],num[s]);
		s=p[top[s]];
	}
	if(dp[s]<dp[t])
		swap(s,t);
	A(d,num[t],num[s]);
}
int calc(int s,int t){
	int v=0;
	while(top[s]!=top[t]){
		if(dp[top[s]]<dp[top[t]])
			swap(s,t);
		v+=Q(num[top[s]],num[s])
		-(U(num[top[s]])
		==U(num[p[top[s]]]));
		s=p[top[s]];
	}
	if(dp[s]<dp[t])
		swap(s,t);
	return v+Q(num[t],num[s]);
}
int main(){
	int m;
	scanf("%d%d",&n,&m);
	for(int i=1;i<=n;++i)
		scanf("%d",w+i);
	for(int i=1;i!=n;++i){
		int u,v;
		scanf("%d%d",&u,&v);
		add(u,v);
	}
	dfs1(1);
	dfs2(top[1]=1);
	build();
	while(m--){
		char d[2];
		int s,t,v;
		scanf("%s%d%d",d,&s,&t);
		if(*d==‘Q‘)
			printf("%d\n",calc(s,t));
		else{
			scanf("%d",&v);
			color(v,s,t);
		}
	}
}