【BZOJ 2243】染色 【树链剖分】

#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
#include<iostream>
#include<cctype>
using namespace std;
#define N 100010
#define g() getchar()
#define d(x) isdigit(x)
#define rep(i,s,t) for(int i=(s);i<=(t);i++)
char ch;template<class T>inline
void F(T& x){
for(ch=g();!d(ch);ch=g());
for(x=0;d(ch);x=x*10+ch-'0',ch=g());
}
int n,m,val
,head
,eid;struct Edge{int to,next;}e[N<<2];
inline void adde(int u,int v){e[eid].to=v;e[eid].next=head[u];head[u]=eid++;}
int size
,dep
,fa
;
void dfs1(int u){
size[u]=1;
for(int i=head[u];~i;i=e[i].next){
if(e[i].to==fa[u])continue;
fa[e[i].to]=u;dep[e[i].to]=dep[u]+1;dfs1(e[i].to);size[u]+=size[e[i].to];
}
}
int tid,top
,pos
;
void dfs2(int x,int topx){
++tid;top[x]=topx;pos[x]=tid;int k=0;
if(size[x]==1)return;
for(int i=head[x];~i;i=e[i].next)
if(dep[e[i].to]>dep[x]&&size[e[i].to]>size[k])
k = e[i].to;
dfs2(k,topx);
for(int i=head[x];~i;i=e[i].next)
if(dep[e[i].to]>dep[x]&&e[i].to!=k)
dfs2(e[i].to,e[i].to);
}
#define ls (u<<1)
#define rs (u<<1|1)
struct SegTree{int l,r,kcnt,lc,rc,cover;}tr[N<<2];
void build(int u,int l,int r){
tr[u].l=l,tr[u].r=r,tr[u].cover=-1;if(l==r)return;
int mid=(l+r)>>1;build(ls,l,mid),build(rs,mid+1,r);
}
inline void pushdown(int u){
tr[ls].lc=tr[ls].rc=tr[rs].lc=tr[rs].rc=tr[ls].cover=tr[rs].cover=tr[u].cover;
tr[ls].kcnt=tr[rs].kcnt=1;tr[u].cover=-1;
}
inline void pushup(int u){
if(tr[ls].rc==tr[rs].lc)tr[u].kcnt=tr[ls].kcnt+tr[rs].kcnt-1;
else tr[u].kcnt=tr[ls].kcnt+tr[rs].kcnt;
tr[u].lc=tr[ls].lc,tr[u].rc=tr[rs].rc;
}
void update(int u,int l,int r,int col){
if(l<=tr[u].l&&tr[u].r<=r){
tr[u].kcnt=1;tr[u].lc=tr[u].rc=col;tr[u].cover=col;return;
}if(~tr[u].cover)pushdown(u);int mid = (tr[u].l+tr[u].r)>>1;
if(l<=mid)update(ls,l,r,col);if(r>mid)update(rs,l,r,col);
pushup(u);
}
int query(int u,int l,int r){
if(l<=tr[u].l&&tr[u].r<=r)return tr[u].kcnt;
int mid=(tr[u].l+tr[u].r)>>1;
if(~tr[u].cover)pushdown(u);
if(r<=mid)return query(ls,l,r);
else if(l>mid)return query(rs,l,r);
else return query(ls,l,r)+query(rs,l,r)+(tr[ls].rc==tr[rs].lc?-1:0);
}
int querycol(int u,int x){
if(tr[u].l==tr[u].r)return tr[u].lc;
if(~tr[u].cover)pushdown(u);
int mid=(tr[u].l+tr[u].r)>>1;
if(x<=mid)return querycol(ls,x);
else return querycol(rs,x);
}
void change(int x,int y,int col){
while(top[x]!=top[y]){
if(dep[top[x]]<dep[top[y]])swap(x,y);
update(1,pos[top[x]],pos[x],col);x=fa[top[x]];
}if(pos[x]>pos[y])swap(x,y);
update(1,pos[x],pos[y],col);
}
int ask(int x,int y){
int ans = 0;
while(top[x]!=top[y]){
if(dep[top[x]]<dep[top[y]])swap(x,y);
ans+=query(1,pos[top[x]],pos[x]);
if(querycol(1,pos[fa[top[x]]])==querycol(1,pos[top[x]]))ans--;
x=fa[top[x]];
}if(pos[x]>pos[y])swap(x,y);
ans+=query(1,pos[x],pos[y]);
return ans;
}
int main(){
int a,b,c;char op[2];memset(head,-1,sizeof(head));
F(n),F(m);rep(i,1,n)F(val[i]);
rep(i,1,n-1){F(a),F(b);adde(a,b);adde(b,a);}
dfs1(1);dfs2(1,1);build(1,1,n);
rep(i,1,n)update(1,pos[i],pos[i],val[i]);
rep(i,1,m){scanf("%s",op);
if(*op=='C'){F(a),F(b),F(c);change(a,b,c);}
else{F(a),F(b);printf("%d\n",ask(a,b));}
}
return 0;
}