BZOJ 2243, 染色

Problem

传送门

Mean

要求编写一个数据结构，维护一颗无根树各节点的颜色以支持查询两点间的路径上颜色段的数量。

Analysis

树链剖分+线段树。
线段树各节点保存区间左端点、右端点的颜色，区间内颜色段数量，以及线段树set操作标记。
合并区间时要注意处理两段点颜色相同的情况。

Code

#include<cstdio>
const int N=100005;
int n,m,x,y,z,ed,dfn,c[N],v[N<<1],nxt[N<<1],g[N],sz[N],f[N],son[N],st[N],top[N],d[N];
char C;
struct Node{
    int l,r,s,f;
}t[N<<2];
void addedge(int x,int y){
    v[++ed]=y;
    nxt[ed]=g[x];
    g[x]=ed;
}
void dfs1(int x){
    sz[x]=1;
    for(int i=g[x];i;i=nxt[i]) if(v[i]!=f[x]){
        f[v[i]]=x,d[v[i]]=d[x]+1;
        dfs1(v[i]);
        sz[x]+=sz[v[i]];
        if(sz[v[i]]>sz[son[x]]) son[x]=v[i];
    }
}
void dfs2(int x,int y){
    st[x]=++dfn,top[x]=y;
    if(son[x]) dfs2(son[x],y);
    for(int i=g[x];i;i=nxt[i]) if(v[i]!=f[x] && v[i]!=son[x]) dfs2(v[i],v[i]);
}
int lca(int x,int y){
    for(;top[x]!=top[y];x=f[top[x]]) if(d[top[x]]<d[top[y]]){int z=x;x=y;y=z;}
    return d[x]<d[y]?x:y;
}
void down(int o){
    int lc=o<<1,rc=o<<1|1;
    t[lc].s=t[rc].s=1;
    t[lc].f=t[rc].f=t[lc].l=t[rc].l=t[lc].r=t[rc].r=t[o].f;
    t[o].f=0;
}
void update(int o,int l,int r,int pl,int pr,int x){
    if(l==pl && r==pr){t[o].l=t[o].r=t[o].f=x;t[o].s=1;}
    else{
        if(t[o].f) down(o);
        int mid=l+r>>1;
        if(mid<pl) update(o<<1|1,mid+1,r,pl,pr,x);
        else if(mid>=pr) update(o<<1,l,mid,pl,pr,x);
        else{update(o<<1,l,mid,pl,mid,x);update(o<<1|1,mid+1,r,mid+1,pr,x);}
        t[o].s=t[o<<1].s+t[o<<1|1].s-(t[o<<1].r==t[o<<1|1].l);
        t[o].l=t[o<<1].l,t[o].r=t[o<<1|1].r;
    }
}
void set(int x,int y,int z){
    while(top[x]!=top[y]){
        update(1,1,n,st[top[x]],st[x],z);
        x=f[top[x]];
    }
    update(1,1,n,st[y],st[x],z);
}
int query(int o,int l,int r,int pl,int pr){
    if(l==pl && r==pr) return t[o].s;
    if(t[o].f) down(o);
    int mid=l+r>>1;
    if(mid<pl) return query(o<<1|1,mid+1,r,pl,pr);
    if(mid>=pr) return query(o<<1,l,mid,pl,pr);
    return query(o<<1,l,mid,pl,mid)+query(o<<1|1,mid+1,r,mid+1,pr)-(t[o<<1].r==t[o<<1|1].l);
}
int col(int o,int l,int r,int p){
    if(l==r) return t[o].l;
    if(t[o].f) down(o);
    int mid=l+r>>1;
    if(p>mid) return col(o<<1|1,mid+1,r,p);
    return col(o<<1,l,mid,p);
}
int calc(int x,int y){
    int sum=0;
    while(top[x]!=top[y]){
        sum+=query(1,1,n,st[top[x]],st[x]);
        if(col(1,1,n,st[top[x]])==col(1,1,n,st[f[top[x]]])) sum--;
        x=f[top[x]];
    }
    return sum+query(1,1,n,st[y],st[x]);
}
int main(){
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++) scanf("%d",&c[i]);
    for(int i=1;i<n;i++){
        scanf("%d%d\n",&x,&y);
        addedge(x,y);
        addedge(y,x);
    }
    dfs1(1);dfs2(1,1);
    for(int i=1;i<=n;i++) update(1,1,n,st[i],st[i],c[i]);
    while(m--){
        if(getchar()=='C'){
            scanf("%d%d%d\n",&x,&y,&z);
            int fa=lca(x,y);
            set(x,fa,z),set(y,fa,z);
        }else{
            scanf("%d%d\n",&x,&y);
            int fa=lca(x,y);
            printf("%d\n",calc(x,fa)+calc(y,fa)-1);
        }
    }
    return 0;
}