BZOJ2243: [SDOI2011]染色 树链剖分

写长程序真闹心。。改了好久

线段树需要存储区间颜色段数，左端点和右端点的颜色。

线段树要注意合并的问题

1.线段树查询跨区间的时候，要注意如果左边子区间的右端点与右边子区间的左端点颜色相同时，颜色段数等于左右区间相加后-1。

2.在树链上也要注意轻链和重链交界处的颜色是否相同，如果颜色相同，颜色段数等于两段链上颜色段数相加后-1

(具体写法参照代码)

因为c的可以等于0，所以标记的原始值要设成-1

#include<cmath>
#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;
const int N=100005;
struct pp
{int l,r,v,col,cor,lazy;}T[N*4];
int n,m;
int a[N];
int to[N*2],nxt[N*2],lj[N],cnt;
void add(int f,int t)
{
	to[++cnt]=t;
	nxt[cnt]=lj[f];
	lj[f]=cnt;
}
int s[N],dep[N],fa[N],son[N];
void dfs1(int x)
{
	s[x]=1;
	dep[x]=dep[fa[x]]+1;
	for(int i=lj[x];i;i=nxt[i])
	if(to[i]!=fa[x])
	{
		fa[to[i]]=x;
		dfs1(to[i]);
		s[x]+=s[to[i]];
		if(s[to[i]]>s[son[x]]) son[x]=to[i];
	}
}
int pos[N],npos[N],top[N],tot;
void dfs2(int x)
{
	pos[x]=++tot;
	npos[tot]=x;
	if(son[fa[x]]==x) top[x]=top[fa[x]];
	else top[x]=x;
	if(son[x]) dfs2(son[x]);
	for(int i=lj[x];i;i=nxt[i])
	if(to[i]!=fa[x]&&to[i]!=son[x]) dfs2(to[i]); 
}
void build(int x,int l,int r)
{
	T[x].l=l,T[x].r=r,T[x].lazy=-1;
	if(l==r)
	{
		T[x].col=a[npos[l]];
		T[x].cor=a[npos[l]];
		T[x].v=1;
		return;
	}
	int mid=(l+r)>>1;
	build(x<<1,l,mid);
	build(x<<1|1,mid+1,r);
	T[x].col=T[x<<1].col;
	T[x].cor=T[x<<1|1].cor;
	T[x].v=T[x<<1].v+T[x<<1|1].v;
	if(T[x<<1].cor==T[x<<1|1].col) T[x].v--;
}
void pushdown(int x)
{
	T[x<<1].col=T[x<<1].cor=T[x].lazy;
	T[x<<1|1].col=T[x<<1|1].cor=T[x].lazy;
	T[x<<1].v=T[x<<1|1].v=1;
	T[x<<1].lazy=T[x<<1|1].lazy=T[x].lazy;
	T[x].lazy=-1;
}
int Query(int x,int l,int r)
{
	if(T[x].l==l&&T[x].r==r) return T[x].v;
	if(T[x].lazy!=-1) pushdown(x);
	int mid=(T[x].l+T[x].r)>>1;
	if(r<=mid)return Query(x<<1,l,r);
	else if(l>mid) return Query(x<<1|1,l,r);
	else
	{
		int p=0;
		if(T[x<<1].cor==T[x<<1|1].col) p=1;
		return Query(x<<1,l,mid)+Query(x<<1|1,mid+1,r)-p;
	}
}
int cQuery(int x,int k)
{
	if(T[x].l==T[x].r) return T[x].cor;
	if(T[x].lazy!=-1) pushdown(x);
	int mid=(T[x].l+T[x].r)>>1;
	if(k<=mid) return cQuery(x<<1,k);
	else return cQuery(x<<1|1,k);
}
int Getnum(int x,int y)
{
	int fx=top[x],fy=top[y],tot=0,cx=-2,cy=-2;
	while(fx!=fy)
	{
		if(dep[fx]<dep[fy])
		{
			tot+=Query(1,pos[fy],pos[y]);
			if(cQuery(1,pos[y])==cy) tot--;
			cy=cQuery(1,pos[fy]);
			y=fa[fy],fy=top[y];
		}
		else
		{
			tot+=Query(1,pos[fx],pos[x]);
			if(cQuery(1,pos[x])==cx) tot--;
			cx=cQuery(1,pos[fx]);
			x=fa[fx],fx=top[x];
		}
	}
	if(dep[x]<dep[y]) tot+=Query(1,pos[x],pos[y]);
	else tot+=Query(1,pos[y],pos[x]);
	if(cQuery(1,pos[x])==cx) tot--;
	if(cQuery(1,pos[y])==cy) tot--;
	return tot;
}
void Modify(int x,int l,int r,int c)
{
	if(T[x].l==l&&T[x].r==r)
	{
		T[x].v=1;
		T[x].col=T[x].cor=c;
		T[x].lazy=c;
		return;
	}
	if(T[x].lazy!=-1) pushdown(x); 
	int mid=(T[x].l+T[x].r)>>1;
	if(r<=mid) Modify(x<<1,l,r,c);
	else if(l>mid) Modify(x<<1|1,l,r,c);
	else Modify(x<<1,l,mid,c),Modify(x<<1|1,mid+1,r,c);
	T[x].col=T[x<<1].col;
	T[x].cor=T[x<<1|1].cor;
	T[x].v=T[x<<1].v+T[x<<1|1].v;
	if(T[x<<1].cor==T[x<<1|1].col) T[x].v--;
}
void Getmodi(int x,int y,int c)
{
	int fx=top[x],fy=top[y];
	while(fx!=fy)
	{
		if(dep[fx]<dep[fy]) swap(fx,fy),swap(x,y);
		Modify(1,pos[fx],pos[x],c);
		x=fa[fx];
		fx=top[x];
	}
	if(dep[x]<dep[y]) swap(x,y);
	Modify(1,pos[y],pos[x],c);
}
int main()
{
	scanf("%d%d",&n,&m);
	for(int i=1;i<=n;i++) scanf("%d",&a[i]);
	int x,y;
	for(int i=1;i<=n-1;i++)
	{
		scanf("%d%d",&x,&y);
		add(x,y),add(y,x);
	}
	dfs1(1),dfs2(1);
	build(1,1,n);
	//for(int i=1;i<=n;i++) cout<<i<<' '<<pos[i]<<endl; 
	for(int i=1;i<=m;i++)
	{
		char o;
		scanf("\n%c%d%d",&o,&x,&y);
		if(o=='Q') printf("%d\n",Getnum(x,y));
		if(o=='C')
		{
			int c;
			scanf("%d",&c);
			Getmodi(x,y,c);
		}
		//for(int i=1;i<=13;i++) cout<<i<<" T[].v="<<T[i].v<<" T[].col="<<T[i].col<<" T[].cor="<<T[i].cor<<endl;
	}
}