bzoj2243[SDOI2011]染色 (树剖)

用线段树维护每段信息,
需要记录每段区间两端的颜色和段数,然后根据情况合并

#include<iostream>
#include<cstdio>
#include<cmath>
#include<cstring>
#include<algorithm>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fod(i,a,b) for(int i=a;i>=b;i--)
using namespace std;
const int N=1e5+10;
int dep[N],s[N],f[N],bl[N],pos[N],size=0,
ls[N<<2],rs[N<<2],ans[N<<2],tag[N<<2],lsgn[N<<2],rsgn[N<<2],
last[N],len=0,n,m,a[N];
struct Edge{int to,next;Edge(int to=0,int next=0):to(to),next(next){}}e[N<<1];
void add_edge(int u,int v){e[++len]=Edge(v,last[u]);last[u]=len;}
void dfs1(int u)
{
    s[u]=1;
    for(int i=last[u];i;i=e[i].next) {
        int id=e[i].to;
        if(id==f[u])continue;
        dep[id]=dep[u]+1;f[id]=u;
        dfs1(id);
        s[u]+=s[id];
    }
}
void dfs2(int u,int chain)
{
    bl[u]=chain;pos[u]=++size;
    int k=0;
    for(int i=last[u];i;i=e[i].next) {
        int id=e[i].to;
        if(dep[id]>dep[u]&&s[id]>s[k])k=id;
    }if(!k)return ;dfs2(k,chain);
    for(int i=last[u];i;i=e[i].next) {
        int id=e[i].to;
        if(dep[id]>dep[u]&&id!=k)dfs2(id,id);
    }
}
void update(int k)
{
    lsgn[k]=lsgn[k<<1];rsgn[k]=rsgn[k<<1|1];
    if(rsgn[k<<1]==lsgn[k<<1|1]) ans[k]=ans[k<<1]+ans[k<<1|1]-1;
    else ans[k]=ans[k<<1]+ans[k<<1|1];
}
void build(int k,int l,int r)
{
    ls[k]=l;rs[k]=r;tag[k]=-1;
    if(l==r)return;
    int mid=(l+r)>>1;
    build(k<<1,l,mid);
    build(k<<1|1,mid+1,r);
}
void pushdown(int k)
{
    ans[k<<1]=ans[k<<1|1]=1;
    tag[k<<1]=tag[k<<1|1]=tag[k];
    lsgn[k<<1]=lsgn[k<<1|1]=rsgn[k<<1|1]=rsgn[k<<1]=tag[k];
    tag[k]=-1;
}
void change(int k,int ql,int qr,int c)
{
    if(ls[k]==ql&&rs[k]==qr) {lsgn[k]=rsgn[k]=c;ans[k]=1;tag[k]=c;return;}
    if(tag[k]!=-1)pushdown(k);
    int mid=(ls[k]+rs[k])>>1;
    if(qr<=mid)change(k<<1,ql,qr,c);
    else if(ql>mid)change(k<<1|1,ql,qr,c);
    else {
        change(k<<1,ql,mid,c);
        change(k<<1|1,mid+1,qr,c);
    }
    update(k);
}
int querysum(int k,int ql,int qr)
{
    if(ql==ls[k]&&qr==rs[k]){return ans[k];}
    if(tag[k]!=-1)pushdown(k);
    int mid=(ls[k]+rs[k])>>1;
    if(qr<=mid)return querysum(k<<1,ql,qr);
    else if(ql>mid)return querysum(k<<1|1,ql,qr);
    else {
        if(rsgn[k<<1]==lsgn[k<<1|1]) return querysum(k<<1,ql,mid)+querysum(k<<1|1,mid+1,qr)-1;
        else return querysum(k<<1,ql,mid)+querysum(k<<1|1,mid+1,qr);
    }
}
int getc(int k,int pos)
{
    if(ls[k]==rs[k]) return lsgn[k];
    if(tag[k]!=-1)pushdown(k);
    int mid=(ls[k]+rs[k])>>1;
    if(pos<=mid)return getc(k<<1,pos);
    else return getc(k<<1|1,pos);
}
int solvesum(int x,int y)
{
    int sum=0;
    while(bl[x]!=bl[y]) {
        if(dep[bl[x]]<dep[bl[y]])swap(x,y);
        sum+=querysum(1,pos[bl[x]],pos[x]);
        if(getc(1,pos[bl[x]])==getc(1,pos[f[bl[x]]]))sum--;
        x=f[bl[x]];
    }if(dep[x]>dep[y])swap(x,y);
    sum+=querysum(1,pos[x],pos[y]);
    return sum;
}
void solvechange(int x,int y,int c)
{
    while(bl[x]!=bl[y]){
        if(dep[bl[x]]<dep[bl[y]])swap(x,y);
        change(1,pos[bl[x]],pos[x],c);
        x=f[bl[x]];
    }if(dep[x]>dep[y])swap(x,y);
    change(1,pos[x],pos[y],c);  
}
int main()
{
    scanf("%d%d",&n,&m);
    fo(i,1,n) scanf("%d",&a[i]);
    for(int from,to,i=1;i<n;i++) {
        scanf("%d%d",&from,&to);
        add_edge(from,to);add_edge(to,from);
    }
    dfs1(1);dfs2(1,1);
    build(1,1,n);
    fo(i,1,n)change(1,pos[i],pos[i],a[i]);
    for(int i=1;i<=m;i++) {
        char ch[10];int x,y,z;
        scanf("%s",ch);
        if(ch[0]=='C'){scanf("%d%d%d",&x,&y,&z);solvechange(x,y,z);}
        if(ch[0]=='Q'){scanf("%d%d",&x,&y);printf("%d\n",solvesum(x,y));}
    }
    return 0;
}