【bzoj 2243】染色（树链剖分）

传送门biu~
树链剖分，在求解的时候注意在两个链的连接处如果颜色相同则答案减一。

#include<bits/stdc++.h>
using namespace std;
int n,m;
struct data{
    int l,r;
    int num,mark;
    int lc,rc;
}tree[400005];
int head[100005],nex[200005],to[200005],tp;
int a[100005],b[100005],siz[100005],dep[100005],fa[100005],son[100005],top[100005],pos[100005],cnt;
inline void add(int x,int y){
    nex[++tp]=head[x];
    head[x]=tp;
    to[tp]=y;
}
void dfs1(int x,int father){
    siz[x]=1;fa[x]=father;dep[x]=dep[father]+1;
    for(int i=head[x];i;i=nex[i]){
        if(to[i]==father)   continue;
        dfs1(to[i],x);
        siz[x]+=siz[to[i]];
        if(siz[to[i]]>siz[son[x]])  son[x]=to[i];
    }
}
void dfs2(int x,int tp){
    top[x]=tp;pos[x]=++cnt;b[cnt]=a[x];
    if(!son[x]) return;
    dfs2(son[x],tp);
    for(int i=head[x];i;i=nex[i]){
        if(to[i]==fa[x] || to[i]==son[x])   continue;
        dfs2(to[i],to[i]);
    }
}
inline void maintain(int x){
    tree[x].num=tree[x<<1].num+tree[x<<1|1].num;
    if(tree[x<<1].rc==tree[x<<1|1].lc)  --tree[x].num;
    tree[x].lc=tree[x<<1].lc;
    tree[x].rc=tree[x<<1|1].rc;
}
inline void pushdown(int x){
    if(!tree[x].mark)   return;
    if(tree[x].l==tree[x].r){
        tree[x].mark=0;
        return;
    }
    tree[x<<1].num=1;tree[x<<1|1].num=1;
    tree[x<<1].lc=tree[x<<1].rc=tree[x<<1|1].lc=tree[x<<1|1].rc=tree[x<<1].mark=tree[x<<1|1].mark=tree[x].mark;
    tree[x].mark=0;
}
void buildtree(int num,int l,int r){
    tree[num].l=l;tree[num].r=r;    
    if(l==r){
        tree[num].num=1;
        tree[num].lc=tree[num].rc=b[l];
        return;
    }
    int mid=l+r>>1;
    buildtree(num<<1,l,mid);
    buildtree(num<<1|1,mid+1,r);
    maintain(num);
}
void Modify(int num,int l,int r,int v){
    pushdown(num);
    if(tree[num].l==l && tree[num].r==r){
        tree[num].num=1;tree[num].mark=v;
        tree[num].lc=tree[num].rc=v;
        return;
    }
    int mid=tree[num].l+tree[num].r>>1;
    if(mid>=r)      Modify(num<<1,l,r,v);
    else if(mid<l)  Modify(num<<1|1,l,r,v);
    else    Modify(num<<1,l,mid,v),Modify(num<<1|1,mid+1,r,v);
    maintain(num);
}
inline void Modify(int x,int y,int v){
    while(top[x]!=top[y]){
        if(dep[top[x]]<dep[top[y]])     swap(x,y);
        Modify(1,pos[top[x]],pos[x],v);
        x=fa[top[x]];
    }
    if(pos[x]>pos[y])   swap(x,y);
    Modify(1,pos[x],pos[y],v);
}
int Query(int num,int l,int r){
    pushdown(num);
    if(tree[num].l==l && tree[num].r==r)    return tree[num].num;
    int mid=tree[num].l+tree[num].r>>1;
    if(mid>=r)      return Query(num<<1,l,r);
    else if(mid<l)  return Query(num<<1|1,l,r);
    else return Query(num<<1,l,mid)+Query(num<<1|1,mid+1,r)-(tree[num<<1].rc==tree[num<<1|1].lc);
}
int Color(int num,int x){
    pushdown(num);
    if(tree[num].l==tree[num].r)    return tree[num].lc;
    int mid=tree[num].l+tree[num].r>>1;
    if(x<=mid)  return Color(num<<1,x);
    else        return Color(num<<1|1,x);
}
inline int Query(int x,int y){
    int re=0;
    while(top[x]!=top[y]){
        if(dep[top[x]]<dep[top[y]])     swap(x,y);
        re+=Query(1,pos[top[x]],pos[x]);
        if(Color(1,pos[top[x]])==Color(1,pos[fa[top[x]]]))  --re;
        x=fa[top[x]];
    }
    if(pos[x]>pos[y])   swap(x,y);
    re+=Query(1,pos[x],pos[y]);
    return re;
}
int main(){
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;++i)   scanf("%d",&a[i]);
    for(int i=1;i<n;++i){
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y);add(y,x);
    }
    dfs1(1,0);dfs2(1,1);
    buildtree(1,1,n);
    while(m--){
        char opt[3];
        scanf("%s",opt);
        if(opt[0]=='C'){
            int x,y,v;
            scanf("%d%d%d",&x,&y,&v);
            Modify(x,y,v);
        }
        else{
            int x,y;
            scanf("%d%d",&x,&y);
            printf("%d\n",Query(x,y));
        }
    }
    return 0;
}