【BZOJ2243】【SDOI2011】染色

  挺裸的一道树链剖分。

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const int maxn=100010;

struct Tree
{
	struct Segment_Tree
	{
		int sum[maxn<<2];
		int lc[maxn<<2],rc[maxn<<2];
		int tag_same[maxn<<2],same[maxn<<2];
		int n;

		#define ls (o<<1)
		#define rs ((o<<1)|1)
		#define mid ((l+r)>>1)

		void Update(int o) {sum[o]=sum[ls]+sum[rs]-(rc[ls]==lc[rs]);lc[o]=lc[ls];rc[o]=rc[rs];}
		
		void Mark_Same(int o,int c) {lc[o]=rc[o]=c;sum[o]=1;tag_same[o]=1;same[o]=c;}

		void Push_Down(int o)
		{
			if (tag_same[o])
			{
				Mark_Same(ls,same[o]);
				Mark_Same(rs,same[o]);
				same[o]=tag_same[o]=0;
			}
		}

		void Build(int o,int l,int r,int *color)
		{
			if (l==r) {sum[o]=1;lc[o]=rc[o]=color[l];return;}
			Build(ls,l,mid,color);
			Build(rs,mid+1,r,color);
			Update(o);
		}

		void Change(int o,int l,int r,int a,int b,int c)
		{
			if (a<=l && r<=b) {Mark_Same(o,c);return;}
			Push_Down(o);
			if (a<=mid) Change(ls,l,mid,a,b,c);
			if (b> mid) Change(rs,mid+1,r,a,b,c);
			Update(o);
		}

		int Query(int o,int l,int r,int a,int b,int &lcolor,int &rcolor)
		{
			if (a<=l && r<=b) {lcolor=lc[o];rcolor=rc[o];return sum[o];}
			Push_Down(o);
			int res=0;
			if (a<=mid && b>mid)
			{
				int ll,lr,rl,rr;
				res+=Query(ls,l,mid,a,b,ll,lr);
				res+=Query(rs,mid+1,r,a,b,rl,rr);
				if (lr==rl) res--;
				lcolor=ll;rcolor=rr;
			}
			else if (a<=mid) res=Query(ls,l,mid,a,b,lcolor,rcolor);else res=Query(rs,mid+1,r,a,b,lcolor,rcolor);
			Update(o);
			return res;
		}
		#undef ls
		#undef rs
		#undef mid
		
		void Build(int *color) {Build(1,1,n,color);}
		void Change(int l,int r,int c) {Change(1,1,n,l,r,c);}
		int Query(int l,int r,int &lcolor,int &rcolor) {return Query(1,1,n,l,r,lcolor,rcolor);}
	}ST;

	int head[maxn],next[maxn<<1],E[maxn<<1],Ecnt;
	int fa[maxn],dep[maxn];
	int top[maxn],wson[maxn];
	int id[maxn],ids;
	int temp[maxn];

	void Add_Edge(int x,int y) {next[++Ecnt]=head[x];head[x]=Ecnt;E[Ecnt]=y;next[++Ecnt]=head[y];head[y]=Ecnt;E[Ecnt]=x;}

	int DFS1(int x,int father,int depth)
	{
		int sons=0,maxsons=0;
		fa[x]=father;dep[x]=depth;
		for (int i=head[x];i;i=next[i]) if (E[i]!=father)
		{
			int tsons=DFS1(E[i],x,depth+1);
			sons+=tsons;
			if (tsons>maxsons) maxsons=tsons,wson[x]=E[i];
		}
		return sons+1;
	}

	void DFS2(int x,int is_wson)
	{
		id[x]=++ids;
		if (is_wson) top[x]=top[fa[x]];else top[x]=x;
		if (wson[x]) DFS2(wson[x],1);
		for (int i=head[x];i;i=next[i]) if (E[i]!=fa[x] && E[i]!=wson[x]) DFS2(E[i],0);
	}

	#define fx (top[x])
	#define fy (top[y])

	void Change(int x,int y,int c)
	{
		while (fx!=fy)
		{
			if (dep[fx]<dep[fy]) swap(x,y);
			ST.Change(id[fx],id[x],c);
			x=fa[fx];
		}
		if (dep[x]>dep[y]) swap(x,y);
		ST.Change(id[x],id[y],c);
	}

	int Query(int x,int y)
	{
		int lc=0,rc=0,tcl=0,tcr=0,res=0;
		while (fx!=fy)
		{
			if (dep[fx]<dep[fy]) swap(x,y),swap(lc,rc);
			res+=ST.Query(id[fx],id[x],tcl,tcr);
			if (tcr==lc) res--;
			lc=tcl;x=fa[fx];
		}
		if (dep[x]>dep[y]) swap(x,y),swap(lc,rc);
		res+=ST.Query(id[x],id[y],tcl,tcr);
		res-=(tcl==lc)+(tcr==rc);
		return res;
	}

	#undef fx
	#undef fy

	void Init(int n,int *color) {DFS1(1,0,1);DFS2(1,0);ST.n=n;for (int i=1;i<=n;i++) temp[id[i]]=color[i];ST.Build(temp);}
}T;

char readchar() {char ch;while ((ch=getchar())!='Q' && ch!='C');return ch;}
void read(int &x) {char ch;while ((ch=getchar())<'0' || ch>'9');x=ch-'0';while ((ch=getchar())<='9' && ch>='0') x=x*10+ch-'0';}
int temp[maxn];
int n,m;

void Init()
{
	read(n),read(m);
	for (int i=1;i<=n;i++) read(temp[i]);
	for (int i=1,x,y;i<n;i++) read(x),read(y),T.Add_Edge(x,y);
	T.Init(n,temp);
}

void Solve()
{
	for (int i=1;i<=m;i++)
	{
		char ch=readchar();
		int a,b,c;
		if (ch=='C') read(a),read(b),read(c),T.Change(a,b,c);
		if (ch=='Q') read(a),read(b),printf("%d\n",T.Query(a,b));
	}
}

int main()
{
	Init();
	Solve();
	return 0;
}