bzoj2243 [SDOI2011]染色（树链剖分,线段树求颜色段数）

树链剖分映射到线段树上，点与点的左右关系，大小关系都是不同的。因此维护ls，rs时要小心。合并答案时看的是top[x]和fa[top[x]]的颜色是否相等，看好是ls还是rs

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define inf 0x3f3f3f3f
#define N 100010
inline int read(){
    int x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-') f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9') x=x*10+ch-'0',ch=getchar();
    return x*f;
}
int n,m,a[N],w[N],h[N],num=0,dep[N],fa[N],size[N],son[N],id[N],dfn=0,top[N];
struct edge{
    int to,next;
}data[N<<1];
struct node{
    int ls,rs,sum,lazy;
    node(){ls=rs=lazy=-1,sum=0;}
    node operator+(node b){
        node res;
        if(ls==-1) res.ls=b.ls;else res.ls=ls;
        if(b.rs==-1) res.rs=rs;else res.rs=b.rs;
        res.sum=sum+b.sum-(rs==b.ls);
        return res;
    }
}tree[N<<2];
void dfs1(int x){
    size[x]=1;
    for(int i=h[x];i;i=data[i].next){
        int y=data[i].to;if(fa[x]==y) continue;
        fa[y]=x;dep[y]=dep[x]+1;dfs1(y);size[x]+=size[y];
        if(size[y]>size[son[x]]) son[x]=y;
    }
}
void dfs2(int x,int tp){
    id[x]=++dfn;top[x]=tp;w[dfn]=a[x];
    if(son[x]) dfs2(son[x],tp);
    for(int i=h[x];i;i=data[i].next){
        int y=data[i].to;
        if(y==son[x]||y==fa[x]) continue;
        dfs2(y,y);
    }
}
void build(int p,int l,int r){
    if(l==r){tree[p].ls=tree[p].rs=w[l];tree[p].sum=1;return;}
    int mid=l+r>>1;
    build(p<<1,l,mid);build(p<<1|1,mid+1,r);
    tree[p]=tree[p<<1]+tree[p<<1|1];
}
inline void pushdown(int p){
    if(tree[p].lazy==-1) return;
    tree[p<<1].lazy=tree[p<<1|1].lazy=tree[p].lazy;
    tree[p<<1].ls=tree[p<<1|1].ls=tree[p].lazy;
    tree[p<<1].rs=tree[p<<1|1].rs=tree[p].lazy;
    tree[p<<1].sum=tree[p<<1|1].sum=1;
    tree[p].lazy=-1;
}
void change(int p,int l,int r,int x,int y,int val){
    if(x<=l&&r<=y){
        tree[p].ls=tree[p].rs=val;tree[p].sum=1;
        tree[p].lazy=val;return;
    }
    int mid=l+r>>1;pushdown(p);
    if(x<=mid) change(p<<1,l,mid,x,y,val);
    if(y>mid) change(p<<1|1,mid+1,r,x,y,val);
    tree[p]=tree[p<<1]+tree[p<<1|1];
}
node query(int p,int l,int r,int x,int y){
    if(x<=l&&r<=y) return tree[p];
    int mid=l+r>>1;pushdown(p);
    node res;
    if(x<=mid) res=query(p<<1,l,mid,x,y)+res;
    if(y>mid) res=res+query(p<<1|1,mid+1,r,x,y);
    return res;
}
void dochange(int x,int y,int val){
    while(top[x]!=top[y]){
        if(dep[top[x]]<dep[top[y]]) swap(x,y);
        change(1,1,n,id[top[x]],id[x],val);
        x=fa[top[x]];
    }
    if(id[x]>id[y]) swap(x,y);change(1,1,n,id[x],id[y],val);
}
int doquery(int x,int y){
    node res1,res2;int res=0;
    while(top[x]!=top[y]){
        if(dep[top[x]]>dep[top[y]]){
            res1=query(1,1,n,id[top[x]],id[x])+res1;
            x=fa[top[x]];
        }
        else{
            res2=query(1,1,n,id[top[y]],id[y])+res2;
            y=fa[top[y]];
        }
    }
    if(id[x]<id[y]){
        node res3=query(1,1,n,id[x],id[y]);
        res=res1.sum+res3.sum+res2.sum-(res1.ls==res3.ls)-(res2.ls==res3.rs);
    }
    else{
        node res3=query(1,1,n,id[y],id[x]);
        res=res1.sum+res3.sum+res2.sum-(res1.ls==res3.rs)-(res2.ls==res3.ls);
    }
    return res;
}
int main(){
//  freopen("a.in","r",stdin);
    n=read();m=read();for(int i=1;i<=n;++i) a[i]=read();
    for(int i=1;i<n;++i){
        int x=read(),y=read();
        data[++num].to=y;data[num].next=h[x];h[x]=num;
        data[++num].to=x;data[num].next=h[y];h[y]=num;
    }
    dep[1]=1;dfs1(1);dfs2(1,1);build(1,1,n);
    while(m--){
        char op[1];scanf("%s",op);int x=read(),y=read();
        if(op[0]=='C') dochange(x,y,read());
        else printf("%d\n",doquery(x,y));
    }
    return 0;
}