bzoj2243-染色(动态树lct)

#include<cstdio>
#include<cstring>
#include<string>
#include<vector>
#include<algorithm>
using namespace std;
const int maxn=100005;
int N,M,C[maxn];
struct lct
{
lct *fa,*son[2];
int rev,c,setc,lc,rc,sum;
};
struct LCT
{
lct data[maxn];
lct *null;
void init(int Size=maxn-1) //初始化
{
null=data; //null指向首元素
for(int i=0;i<=Size;i++)
{
data[i].son[0]=data[i].son[1]=data[i].fa=null;
data[i].rev=0;
data[i].c=data[i].lc=data[i].rc=C[i];
data[i].setc=-1;
data[i].sum=1;
}
null->c=null->lc=null->rc=-1;
null->sum=0;
}
void push_rev(lct* x)
{
if(x==null) return;
x->rev=!x->rev;
swap(x->son[0],x->son[1]);
swap(x->lc,x->rc);
}
void push_setc(lct* x,int setc)
{
if(x==null) return;
x->c=setc;
x->setc=setc;
x->lc=setc;
x->rc=setc;
x->sum=1;
}
void pushdown(lct* x)
{
if(x->setc!=-1)
{
push_setc(x->son[0],x->setc);
push_setc(x->son[1],x->setc);
x->setc=-1;
}
if(x->rev)
{
push_rev(x->son[0]);
push_rev(x->son[1]);
x->rev=0;
}
}
void pushup(lct* x)
{
if(x==null) return;
x->sum=1;
if(x->son[0]!=null)
{
x->sum+=x->son[0]->sum;
if(x->son[0]->rc==x->c) x->sum--;
}
if(x->son[1]!=null)
{
x->sum+=x->son[1]->sum;
if(x->son[1]->lc==x->c) x->sum--;
}
x->lc=x->c;
if(x->son[0]!=null) x->lc=x->son[0]->lc;
x->rc=x->c;
if(x->son[1]!=null) x->rc=x->son[1]->rc;
}
bool Same(lct* x,lct* &y) //判断x和x的父亲是否在同一树里
{
return (y=x->fa)!=null&&(y->son[0]==x||y->son[1]==x);
}
void Rotate(lct* x,int d) //翻转
{
lct* y=x->fa; //x的父亲
y->son[d^1]=x->son[d];
if(x->son[d]!=null) x->son[d]->fa=y; //x的子节点的父亲指向y
x->fa=y->fa;  //连接
if(y->fa->son[0]==y) x->fa->son[0]=x;
else if(y->fa->son[1]==y) x->fa->son[1]=x;
x->son[d]=y;
y->fa=x;
}
void Splay(lct* x)
{
pushdown(x); //清除标记
lct* y;
while(Same(x,y)) //没有到树的最顶点
{
pushdown(y);
pushdown(x);
Rotate(x,y->son[0]==x); //翻转
pushup(y);
pushup(x);
}
}
lct* Access(lct* u)  //打通路径,返回的是根
{
lct *v=null;
for(;u!=null;u=u->fa)
{
Splay(u);
u->son[1]=v;
pushup(v=u);
}
return v;
}
lct* GetRoot(lct* x) //得到根
{
for(x=Access(x);pushdown(x),x->son[0]!=null;x=x->son[0]) pushup(x);
return x;
}
void MakeRoot(lct* x) //使x成为根
{
Access(x);
Splay(x);
push_rev(x);
}
void Link(lct* x,lct* y) //连接两个点
{
MakeRoot(x);
x->fa=y;
Access(x);
}
void Cut(lct* x,lct* y) //断开两个点
{
MakeRoot(x);
Access(y);
Splay(y);
y->son[0]->fa=null;
y->son[0]=null;
}
void SetC(lct* x,lct* y,int setc)
{
MakeRoot(x);
push_setc(Access(y),setc);
}
int Query(lct* x,lct* y)
{
MakeRoot(x);
Access(y);
Splay(y);
return y->sum;
}
}A;
int main()
{
scanf("%d%d",&N,&M);
for(int i=1;i<=N;i++) scanf("%d",&C[i]);
A.init(N);
int x,y;
for(int i=1;i<N;i++)
{
scanf("%d%d",&x,&y);
A.Link(A.data+x,A.data+y);
}
char op[3];
int c;
while(M--)
{
scanf("%s%d%d",op,&x,&y);
if(op[0]=='C') scanf("%d",&c);
if(op[0]=='C') A.SetC(A.data+x,A.data+y,c);
else printf("%d\n",A.Query(A.data+x,A.data+y));
}
return 0;
}