【bzoj 2011】染色 树剖+线段树

今天考试的第4题，我的心情有点崩溃，只写了半小时，最后wa了，可是下来百思不得其解，结果结果。。。。我把一个地方的右儿子u<<1|1写成了u<<1主要是写的时候太激动，主要是看错时间，以为只有半小时就结束了结果，还有1个半小时

思路没什么好讲的，闭着眼都知道是树剖主要是细节方面，向上移动的时候判断颜色是否相同就好了，还有，懒惰标记可能是0，所以标记成-1（这个也很重要）

#include<cstdio>
#include<cstring>
#include<iostream>
#define ls u<<1,l,mid
#define rs u<<1|1,mid+1,r
#define maxn 200020
using namespace std;
void init(){
	freopen("paint.in","r",stdin);
	freopen("paint.out","w",stdout);
}
int n,m,top[maxn],f[maxn],size[maxn],tot=1,cnt,dfn[maxn],son[maxn],head[maxn*3],col[maxn],h[maxn],nu[maxn];
struct edge{
	int v,next;
}e[maxn*3];
void adde(int a,int b){
	e[tot].v=b;
	e[tot].next=head[a];
	head[a]=tot++;
}
void dfs1(int u,int fa){
	size[u]=1;h[u]=h[fa]+1;f[u]=fa;
	for(int i=head[u];~i;i=e[i].next){
		int v=e[i].v;if(v==fa)continue;
		dfs1(v,u);
		size[u]+=size[v];
		if(size[v]>size[son[u]])son[u]=v;
	}
}
void dfs2(int u,int fa,int ff){
	dfn[u]=++cnt;nu[cnt]=u;top[u]=ff;
	if(son[u])dfs2(son[u],u,ff);
	for(int i=head[u];~i;i=e[i].next){
		int v=e[i].v;if(v==fa||v==son[u])continue;
		dfs2(v,u,v);
	}
}
struct node{
	int l,r,sum,lc,rc,lazy;
}nod[maxn*4];
void push_up(int u){
	nod[u].lc=nod[u<<1].lc ,nod[u].rc=nod[u<<1|1].rc ;//我擦我擦我擦我擦我擦我擦我擦 
	nod[u].sum=nod[u<<1].sum+nod[u<<1|1].sum;
	if(nod[u<<1].rc==nod[u<<1|1].lc)nod[u].sum--;
}
void push_down(int u){
	int ll=u<<1,rr=u<<1|1;
	if(nod[u].lazy==-1)return ;
	int add=nod[u].lazy;
	nod[ll].lazy=nod[rr].lazy=nod[ll].lc=nod[ll].rc=nod[rr].lc=nod[rr].rc=add;
	nod[ll].sum=nod[rr].sum=1;
	nod[u].lazy=-1;
}
void build(int u,int l,int r){
	nod[u].l=l,nod[u].r=r,nod[u].lazy=-1;
	if(l==r){
		nod[u].sum=1;nod[u].lc=nod[u].rc=col[nu[l]];
		return;
	}
	int mid=l+r>>1;
	build(ls);build(rs);
	push_up(u);
}
void update(int u,int l,int r,int x,int y,int add){
	if(l==x&&y==r){
		nod[u].sum=1;nod[u].lc=nod[u].rc=nod[u].lazy=add;
		
		return;
	}
	push_down(u);
	int mid=l+r>>1;
	if(x>mid)update(rs,x,y,add);
	else if(y<=mid)update(ls,x,y,add);
	else update(ls,x,mid,add),update(rs,mid+1,y,add);
	push_up(u);
}
int query(int u,int l,int r,int x,int y){
	if(x==l&&y==r){
		return nod[u].sum;
	}
	push_down(u);
	int mid=l+r>>1;
	if(x>mid)return query(rs,x,y);
	else if(y<=mid)return query(ls,x,y);
	else {
		int ans=query(ls,x,mid)+query(rs,mid+1,y);
		if(nod[u<<1].rc==nod[u<<1|1].lc)ans--;
		return ans;
	}
}
int query_colour(int u,int l,int r,int x){
	if(l==r)return nod[u].lc;
	push_down(u);
	int mid=l+r>>1;
	if(x>mid)return query_colour(rs,x);
	else return query_colour(ls,x);
}

void change(int a,int b,int add){
	while(top[a]!=top[b]){
		if(h[top[a]]>h[top[b]])swap(a,b);
		update(1,1,cnt,dfn[top[b]],dfn[b],add);
		b=f[top[b]];
	}
	if(h[a]>h[b])swap(a,b);
	update(1,1,cnt,dfn[a],dfn[b],add);
}


int QUERY(int a,int b){
	int ans=0,nc,fc;
	while(top[a]!=top[b]){
		if(h[top[a]]>h[top[b]])swap(a,b);
		ans+=query(1,1,cnt,dfn[top[b]],dfn[b]);
		nc=query_colour(1,1,cnt,dfn[top[b]]);
		fc=query_colour(1,1,cnt,dfn[f[top[b]]]);
		b=f[top[b]];
		if(nc==fc)ans--;
	}
	if(h[a]>h[b])swap(a,b);
	ans+=query(1,1,cnt,dfn[a],dfn[b]);
	return ans==0?1:ans;
}

int main(){
	//init();
	scanf("%d%d",&n,&m);
	memset(head,-1,sizeof(head));
	for(int i=1;i<=n;i++)scanf("%d",col+i);
	for(int a,b,i=1;i<n;i++){
		scanf("%d%d",&a,&b);
		adde(a,b);adde(b,a);
	}
	dfs1(1,1);dfs2(1,1,1);
	build(1,1,cnt);
	/*for(int i=1;i<=n;i++){
		printf("%d %d\n",query_colour(1,1,cnt,dfn[i]),query(1,1,cnt,dfn[i],dfn[i]));
	}*/
	char s[5];
	int a,b,c;
	/*while(scanf("%d%d",&a,&b)!=EOF){
		printf("%d\n",query(1,1,cnt,dfn[a],dfn[b]));
	}*/
	while(m--){
		scanf("%s%d%d",s,&a,&b);
		if(s[0]=='Q')printf("%d\n",QUERY(a,b));
		else {
			scanf("%d",&c);
			change(a,b,c);
		}
	}
	return 0;
}