P2486 [SDOI2011]染色 树剖

刚刚还给征宇讲：我他妈最讨厌这样贼长的代码了，思路很清晰，样例数据也能过，但是就是wa

 

幸福来得有点猝不及防~

一发AC

 

 

#include<bits/stdc++.h>
using namespace std;
const int M=1e5+50;
struct Edge{int to,next;}e[M<<1];
struct Node{int l,r,sum,lc,rc,v;}t[M<<2];
int n,m,cnt,head[M<<1],a[M];
int tot,dep[M],siz[M],son[M],fa[M],id[M],rk[M],top[M];

void add(int x,int y);
void dfs1(int x,int f);
void dfs2(int x,int tp);
void push_up(int k);
void build(int k,int l,int r);
void down(int k);
void update(int k,int l,int r,int c);
void updates();
int ask(int k,int l,int r);
int ask1(int k,int in);
void query();

int main(){
	memset(head,-1,sizeof(head));cnt=1;
	scanf("%d%d",&n,&m);
	for(int i=1;i<=n;i++)scanf("%d",&a[i]);
	for(int x,y,i=1;i<n;i++)scanf("%d%d",&x,&y),add(x,y);
	dfs1(1,0);
	dfs2(1,1);
	build(1,1,n);
	char op[3];
	while(m--){
		scanf("%s",op);
		if(op[0]=='C')updates();
		else query();
	}
	return 0;
}


void add(int x,int y){
	e[cnt].to=y;
	e[cnt].next=head[x];
	head[x]=cnt++;
	e[cnt].to=x;
	e[cnt].next=head[y];
	head[y]=cnt++;
}
void dfs1(int x,int f){
	fa[x]=f;
	dep[x]=dep[f]+1;
	siz[x]=1;
	for(int to,i=head[x];~i;i=e[i].next){
		to=e[i].to;
		if(to==f)continue;
		dfs1(to,x);
		siz[x]+=siz[to];
		if(siz[ son[x] ]<siz[to]) son[x]=to;
	}
}
void dfs2(int x,int tp){
	top[x]=tp;
	id[x]=++tot;
	rk[tot]=a[x];
	if(!son[x])return ;
	dfs2(son[x],tp);
	for(int i=head[x];~i;i=e[i].next)
		if(!id[e[i].to])
			dfs2(e[i].to,e[i].to);
}
void push_up(int k){
	t[k].sum=t[k*2].sum+t[k*2+1].sum;
	t[k].lc=t[k*2].lc;
	t[k].rc=t[k*2+1].rc;
	if(t[k*2].rc==t[k*2+1].lc)t[k].sum--;
}
void build(int k,int l,int r){
	t[k].l=l;t[k].r=r;t[k].v=0;
	if(l==r){
		t[k].lc=t[k].rc=rk[l];t[k].sum=1;
		return ;
	}
	int mid=(l+r)>>1;
	build(k*2,l,mid);
	build(k*2+1,mid+1,r);
	push_up(k);
}

void down(int k){
	if(t[k].v){
		t[k*2].lc=t[k*2].rc=t[k*2].v=t[k].v;
		t[k*2+1].lc=t[k*2+1].rc=t[k*2+1].v=t[k].v;
		t[k*2].sum=t[k*2+1].sum=1;
		t[k].v=0;
	}
}
void update(int k,int l,int r,int c){
	if(l<=t[k].l&&t[k].r<=r){
		t[k].sum=1;
		t[k].lc=t[k].rc=t[k].v=c;
		return ;
	}
	int  mid=(t[k].l+t[k].r)>>1;
	down(k);
	if(l<=mid)update(k*2,l,r,c);
	if(mid<r)update(k*2+1,l,r,c);
	push_up(k);
}
void updates(){
	int x,y,c;scanf("%d%d%d",&x,&y,&c);
	while(top[x]!=top[y]){
		if(dep[top[x]]<dep[top[y]])swap(x,y);
		update(1,id[top[x]],id[x],c);
		x=fa[top[x]];
	}
	if(id[x]>id[y])swap(x,y);
	update(1,id[x],id[y],c);
}
int ask(int k,int l,int r){
	if(l<=t[k].l&&t[k].r<=r)return t[k].sum;
	down(k);
	int res=0,mid=(t[k].l+t[k].r)>>1;
	if(l<=mid)res+=ask(k*2,l,r);
	if(mid<r)res+=ask(k*2+1,l,r);
	if(l<=mid&&mid<r&&t[k*2].rc==t[k*2+1].lc)res--;
	return res;
}
int ask1(int k,int in){
	if(t[k].l==t[k].r)return t[k].lc;
	int  mid=(t[k].l+t[k].r)>>1;
	down(k);
	if(in<=mid) return ask1(k*2,in);
	return ask1(k*2+1,in);
}

void query(){
	int x,y,res=0,c1,c2;
	scanf("%d%d",&x,&y);
	while(top[x]!=top[y]){
		if(dep[top[x]]<dep[top[y]])swap(x,y);
		res+=ask(1,id[top[x]],id[x]);
		if(ask1(1,id[fa[top[x]]])==ask1(1,id[top[x]]))res--;
		x=fa[top[x]];
	}
	if(id[x]>id[y])swap(x,y);
	res+=ask(1,id[x],id[y]);
	printf("%d\n",res);	
}