【BZOJ2243】【SDOI2011】染色

题意见试题传送门

解题思路：显然是题树剖题。

考虑用线段树维护区间端点颜色与颜色数，这样就可以方便的合并，注意查询的时候对端点的特殊处理即可。

时间效率最高为\( O (m \log^{2} n) \).(BZOJ 上 4072ms)

#include <stdio.h>
#define MN 100005
#define Mn (1<<17)
#define ls(x) (x<<1)
#define rs(x) (x<<1|1)
#define mid (l+r>>1)
int lc[Mn<<1],rc[Mn<<1],sum[Mn<<1],mark[Mn<<1];
int n,q,col[MN],siz[MN],fa[MN],son[MN],dep[MN],top[MN],pos[MN],head[MN],cnt,dfsn,rev[MN];
int to[MN<<1],nxt[MN<<1];
inline int in(){
    int x=0;bool f=0; char ch=getchar();
    while(ch<‘0‘||ch>‘9‘) f=ch==‘-‘,ch=getchar();
    while(ch>=‘0‘&&ch<=‘9‘) x=(x<<3)+(x<<1)+ch-‘0‘,ch=getchar();
    return f?-x:x;
}
inline void swp(int &a,int &b){a^=b^=a^=b;}
inline void ins(int x,int y){to[++cnt]=y,nxt[cnt]=head[x],head[x]=cnt;}
inline void dfs1(int u,int f,int d){
    fa[u]=f,dep[u]=d,siz[u]=1;
    for (register int i=head[u]; i; i=nxt[i])
        if (to[i]!=f){
            dfs1(to[i],u,d+1);siz[u]+=siz[to[i]];
            if (siz[to[i]]>siz[son[u]]) son[u]=to[i];
        }
}
inline void dfs2(int u,int tp){
    top[u]=tp;pos[u]=(++dfsn);rev[dfsn]=u;if (son[u]) dfs2(son[u],tp);
    for (register int i=head[u]; i; i=nxt[i])
        if (to[i]!=fa[u]&&to[i]!=son[u]) dfs2(to[i],to[i]);
}
inline void pushdown(int k){
    mark[ls(k)]=mark[rs(k)]=lc[ls(k)]=rc[ls(k)]=lc[rs(k)]=rc[rs(k)]=mark[k];
    sum[ls(k)]=sum[rs(k)]=1;mark[k]=0;
}
inline void combine(int k){
    sum[k]=rc[ls(k)]==lc[rs(k)]?sum[ls(k)]+sum[rs(k)]-1:sum[ls(k)]+sum[rs(k)];
    lc[k]=lc[ls(k)],rc[k]=rc[rs(k)];
}
inline void build(int k,int l,int r){
    mark[k]=0;
    if (l==r){sum[k]=1,lc[k]=rc[k]=col[rev[l]];return;}
    build(ls(k),l,mid);build(rs(k),mid+1,r);combine(k);
}
inline void A(int l,int r,int a,int b,int k,int col){
    if (a<=l&&r<=b){mark[k]=lc[k]=rc[k]=col;sum[k]=1;return;}
    if (mark[k]) pushdown(k);if (a<=mid) A(l,mid,a,b,ls(k),col);
    if (b>mid) A(mid+1,r,a,b,rs(k),col);combine(k);
}
inline int Qs(int l,int r,int a,int b,int k){
    if (l==a&&r==b) return sum[k];if (mark[k]) pushdown(k);
    if (b<=mid) return Qs(l,mid,a,b,ls(k));
    if (a>mid) return Qs(mid+1,r,a,b,rs(k));
    return Qs(l,mid,a,mid,ls(k))+Qs(mid+1,r,mid+1,b,rs(k))-(rc[ls(k)]==lc[rs(k)]);
}
inline int Qc(int l,int r,int x,int k){
    if (l==r) return lc[k];if (mark[k]) pushdown(k);
    if (x<=mid) return Qc(l,mid,x,ls(k));return Qc(mid+1,r,x,rs(k));
}
inline void update(int x,int y,int cl){
    while (top[x]!=top[y]){
        if (dep[top[x]]<dep[top[y]]) swp(x,y);
        A(1,n,pos[top[x]],pos[x],1,cl);x=fa[top[x]];
    }if (dep[x]>dep[y]) swp(x,y);A(1,n,pos[x],pos[y],1,cl);
}
inline int query(int x,int y){
    register int res=0;
    while(top[x]!=top[y]){
        if (dep[top[x]]<dep[top[y]]) swp(x,y);
        res+=Qs(1,n,pos[top[x]],pos[x],1)-(Qc(1,n,pos[top[x]],1)==Qc(1,n,pos[fa[top[x]]],1));
        x=fa[top[x]];
    }if (dep[x]>dep[y]) swp(x,y);res+=Qs(1,n,pos[x],pos[y],1);return res;
}
void init(){
    n=in(),q=in();for (int i=1; i<=n; ++i) col[i]=in();
    for (register int i=1; i<n; ++i){
        register int x=in(),y=in();
        ins(x,y); ins(y,x);
    }dfs1(1,1,1);dfs2(1,1);build(1,1,n);
}
void solve(){
    while(q--){
        register char ch=getchar();while (ch!=‘C‘&&ch!=‘Q‘) ch=getchar();
        register int x=in(),y=in();
        if (ch==‘C‘) update(x,y,in());
        else printf("%d\n",query(x,y));
    }
}
int main(){init(); solve(); return 0;}