bzoj2243 染色（树链剖分+线段树）

题目链接：bzoj2243 染色
题意：一棵树，有点权，覆盖链上的颜色，询问链上颜色的段数。

序列的颜色段数很好求，只要维护区间颜色段数，以及最左端与最右端的颜色种类即可。
考虑树上操作，可以树链剖分将树上操作转化为链上操作。
修改很好做查询需注意。
链上版本可以拆为两个：x往上跳，y往上跳，可以自己定义序列方向为x->y，那么只需要合并时候注意一下就好。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn=1e5+7;

struct Edge{
    int v,next;
}edge[maxn<<1];

int head[maxn],top;
void init(){
    top=0;
    memset(head,-1,sizeof(head));
}

void add(int u,int v){
    edge[top].v=v;
    edge[top].next=head[u];
    head[u]=top++;
}

int dep[maxn],val[maxn],fa[maxn],dfn[maxn],num,size1[maxn],son[maxn],topp[maxn],W[maxn];
bool vis[maxn];

void dfs1(int u){
    vis[u]=1;
    size1[u]=1;
    int v;
    int maxx=-1;
    for(int i=head[u];i!=-1;i=edge[i].next){
        v=edge[i].v;
        if(vis[v]) continue;
        dep[v]=dep[u]+1;
        fa[v]=u;
        dfs1(v);
        size1[u]+=size1[v];
        if(size1[v]>maxx){
            son[u]=v;
            maxx=size1[v];
        }
    }
}

void dfs2(int u,int last){
    topp[u]=last;
    dfn[u]=++num;
    val[num]=W[u];
    if(!son[u]) return ;
    dfs2(son[u],last);
    int v;
    for(int i=head[u];i!=-1;i=edge[i].next){
        v=edge[i].v;
        if(dfn[v]||v==son[u]) continue;
        dfs2(v,v);
    }
}
struct Tree{
    int dat,lc,rc;
    void init(){
        dat=0;
        lc=rc=-1;
    }
}a[maxn<<2|1];
int lazy[maxn<<2|1];
void pushup(int k){
    a[k].dat=a[k<<1].dat+a[k<<1|1].dat;
    a[k].lc=a[k<<1].lc,a[k].rc=a[k<<1|1].rc;
    if(a[k<<1].rc==a[k<<1|1].lc) --a[k].dat;
}

void pushdown(int k){
    if(lazy[k]!=-1){
        lazy[k<<1]=lazy[k<<1|1]=lazy[k];
        a[k<<1].dat=a[k<<1|1].dat=1;
        a[k<<1].lc=a[k<<1].rc=a[k<<1|1].lc=a[k<<1|1].rc=lazy[k];
        lazy[k]=-1;
    }
}

void build(int l,int r,int k){
    lazy[k]=-1;
    if(l==r){
        a[k].dat=1;
        a[k].lc=a[k].rc=val[l];
        return ;
    }
    int mid=(l+r)>>1;
    build(l,mid,k<<1);
    build(mid+1,r,k<<1|1);
    pushup(k);
}

void updata(int l,int r,int k,int L,int R,int val){
    if(l>=L&&r<=R){
        lazy[k]=a[k].lc=a[k].rc=val;
        a[k].dat=1;
        return ;
    }
    pushdown(k);
    int mid=(l+r)>>1;
    if(L<=mid) updata(l,mid,k<<1,L,R,val);
    if(R>mid) updata(mid+1,r,k<<1|1,L,R,val);
    pushup(k);
}

Tree myfind(int l,int r,int k,int L,int R){
    if(l>=L&&r<=R) return a[k];
    pushdown(k);
    int mid=(l+r)>>1;
    Tree x,y,z;
    if(L<=mid) z=x=myfind(l,mid,k<<1,L,R);
    if(R>mid) z=y=myfind(mid+1,r,k<<1|1,L,R);

    if(L<=mid&&R>mid){
        z.dat=x.dat+y.dat;
        z.lc=x.lc,z.rc=y.rc;
        if(x.rc==y.lc) --z.dat;
    }
    pushup(k);
    return z;
}

int n;
void updatachain(int x,int y,int val){
    while(topp[x]!=topp[y]){
        if(dep[topp[x]]<dep[topp[y]]) swap(x,y);
        updata(1,n,1,dfn[topp[x]],dfn[x],val);
        x=fa[topp[x]];
    }
    if(dep[x]>dep[y]) swap(x,y);
    updata(1,n,1,dfn[x],dfn[y],val);
}
void out(Tree a){
    cout<<a.dat<<" "<<a.lc<<" "<<a.rc<<endl;
}
int myfindchain(int x,int y){
    Tree resu,resv,u,v;
    int res=0;
    resu.init(),resv.init();
    while(topp[x]!=topp[y]){
        if(dep[topp[x]]>dep[topp[y]]){
            u=myfind(1,n,1,dfn[topp[x]],dfn[x]),x=fa[topp[x]];
            //out(u);
            if(resu.dat){
                //out(resu);
                resu.dat+=u.dat;
                if(resu.rc==u.rc) --resu.dat;
                resu.rc=u.lc;
            }
            else resu=u,swap(resu.lc,resu.rc);
        }
        else{
            v=myfind(1,n,1,dfn[topp[y]],dfn[y]),y=fa[topp[y]];
            //out(v);
            if(resv.dat){
                //out(resv);
                resv.dat+=v.dat;
                if(resv.lc==v.rc) --resv.dat;
                resv.lc=v.lc;
            }
            else resv=v;
        }
    }
    if(dep[x]<dep[y]){
        u=myfind(1,n,1,dfn[x],dfn[y]);
        //out(u);
        res=resu.dat+resv.dat+u.dat;
        if(resu.rc==u.lc) --res;
        if(u.rc==resv.lc) --res;
    }
    else{
        v=myfind(1,n,1,dfn[y],dfn[x]);
        //out(v);
        res=resu.dat+resv.dat+v.dat;
        if(resu.rc==v.rc) --res;
        if(v.lc==resv.lc) --res;
    }
    return res;
}
char s[9];
int main(){
    //freopen("in.txt","r",stdin);
    //freopen("out1.txt","w",stdout);
    init();
    int q;
    scanf("%d%d",&n,&q);
    for(int i=1;i<=n;++i) scanf("%d",&W[i]);
    int u,v,w,id;
    for(int i=1;i<n;++i){
        scanf("%d%d",&u,&v);
        add(u,v),add(v,u);
    }
    dep[1]=1;
    dfs1(1);
    dfs2(1,1);
    build(1,n,1);
    while(q--){
        scanf("%s%d%d",s,&u,&v);
        if(s[0]=='C'){
            scanf("%d",&w);
            updatachain(u,v,w);
        }
        else printf("%d\n",myfindchain(u,v));
    }
    return 0;
}