BZOJ：2243: [SDOI2011]染色

题解：

树剖，线段树维护区间颜色段数

记录两端点的颜色，做到O(1)合并

问题：

　　非递归建树实现

#include<iostream>
#include<cstdio>
#include<cstring>
using namespace std;
const int maxn=100009;

int n,m;
int incolor[maxn];

int cntedge;
int head[maxn];
int to[maxn<<1],nex[maxn<<1];
int addedge(int x,int y){
    nex[++cntedge]=head[x];
    to[cntedge]=y;
    head[x]=cntedge;
}

int father[maxn],depth[maxn],siz[maxn];
int heavyson[maxn],top[maxn],idx[maxn],ref[maxn];

int dfs1(int now,int fa){
    father[now]=fa;
    depth[now]=depth[fa]+1;
    siz[now]=1;
    for(int i=head[now];i;i=nex[i]){
        if(to[i]==fa)continue;
        dfs1(to[i],now);
        siz[now]+=siz[to[i]];
        if(siz[to[i]]>siz[heavyson[now]]){
            heavyson[now]=to[i];
        }
    }
}
int temp;
int dfs2(int now,int toppoint){
    top[now]=toppoint;
    idx[now]=(++temp);
    ref[temp]=now;
    if(!heavyson[now])return 0;
    dfs2(heavyson[now],toppoint);
    for(int i=head[now];i;i=nex[i]){
        if(to[i]==father[now])continue;
        if(to[i]==heavyson[now])continue;
        dfs2(to[i],to[i]);
    }
}

struct SegmentTree{
    int l,r;
    int lcolor,rcolor;
    int setcolor,sum;
}tree[maxn<<2];
int pushup(int now){
    tree[now].lcolor=tree[now<<1].lcolor;
    tree[now].rcolor=tree[now<<1|1].rcolor;
    if(tree[now<<1].rcolor==tree[now<<1|1].lcolor){
        tree[now].sum=tree[now<<1].sum+tree[now<<1|1].sum-1;
    }else{
        tree[now].sum=tree[now<<1].sum+tree[now<<1|1].sum;
    }
}
int pushdown(int now){
    if(tree[now].setcolor!=-1){
        tree[now<<1].setcolor=tree[now<<1].lcolor=tree[now<<1].rcolor=tree[now].setcolor;
        tree[now<<1|1].setcolor=tree[now<<1|1].lcolor=tree[now<<1|1].rcolor=tree[now].setcolor;
        tree[now<<1].sum=tree[now<<1|1].sum=1;
        tree[now].setcolor=-1;
    }
}

int Build(int now,int l,int r){
    tree[now].l=l;tree[now].r=r;
    tree[now].setcolor=-1;
    if(l==r){
        tree[now].lcolor=tree[now].rcolor=incolor[ref[l]];
        tree[now].sum=1;
        return 0;
    }
    int mid=(l+r)>>1;
    Build(now<<1,l,mid);
    Build(now<<1|1,mid+1,r);
    pushup(now);
}

int Updata(int now,int ll,int rr,int x){
    if(tree[now].l>=ll&&tree[now].r<=rr){
        tree[now].setcolor=tree[now].lcolor=tree[now].rcolor=x;
        tree[now].sum=1;
        return 0;
    }
    int mid=(tree[now].l+tree[now].r)>>1;
    pushdown(now);
    if(ll<=mid)Updata(now<<1,ll,rr,x);
    if(rr>mid)Updata(now<<1|1,ll,rr,x);
    pushup(now);
}

int Querysum(int now,int ll,int rr){
    if(tree[now].l>=ll&&tree[now].r<=rr){
        return tree[now].sum;
    }
    int mid=(tree[now].l+tree[now].r)>>1;
    pushdown(now);
    if(rr<=mid)return Querysum(now<<1,ll,rr);
    else if(ll>mid)return Querysum(now<<1|1,ll,rr);
    else if(tree[now<<1].rcolor==tree[now<<1|1].lcolor)return Querysum(now<<1,ll,rr)+Querysum(now<<1|1,ll,rr)-1;
    else return Querysum(now<<1,ll,rr)+Querysum(now<<1|1,ll,rr);
}

int Querycolor(int now,int p){
    if(tree[now].l==tree[now].r){
        return tree[now].lcolor;
    }
    int mid=(tree[now].l+tree[now].r)>>1;
    pushdown(now);
    if(p<=mid)return Querycolor(now<<1,p);
    else return Querycolor(now<<1|1,p);
}

int Change(int u,int v,int x){
    int tu=top[u];
    int tv=top[v];
    while(tu!=tv){
        if(depth[tu]<depth[tv]){
            swap(tu,tv);swap(u,v);
        }
        Updata(1,idx[tu],idx[u],x);
        u=father[tu];tu=top[u];
    }
    if(depth[u]>depth[v])swap(u,v);
    Updata(1,idx[u],idx[v],x);
}

int Getans(int u,int v){
    int ret=0;
    int tu=top[u];
    int tv=top[v];
    while(tu!=tv){
        if(depth[tu]<depth[tv]){
            swap(tu,tv);swap(u,v);
        }
        ret=ret+Querysum(1,idx[tu],idx[u]);
        u=father[tu];
        if(Querycolor(1,idx[u])==Querycolor(1,idx[tu]))--ret;
        tu=top[u];
    }
    if(depth[u]>depth[v])swap(u,v);
    ret=ret+Querysum(1,idx[u],idx[v]);
    return ret;
}

int minit(){
    temp=cntedge=0;
    memset(heavyson,0,sizeof(heavyson));
    memset(head,0,sizeof(head));
}

int main(){
    scanf("%d%d",&n,&m);
    minit();
    
    for(int i=1;i<=n;++i)scanf("%d",&incolor[i]);
    for(int i=1;i<=n-1;++i){
        int x,y;
        scanf("%d%d",&x,&y);
        addedge(x,y);addedge(y,x);
    }
    dfs1(1,0);
    dfs2(1,1);
    Build(1,1,n);
    
    char opty[4];
    while(m--){
        int x,y,z;
        scanf("%s",opty);
        if(opty[0]=='C'){
            scanf("%d%d%d",&x,&y,&z);
            Change(x,y,z);
        }else{
            scanf("%d%d",&x,&y);
            printf("%d\n",Getans(x,y));
        }
    }
    return 0;
}