BZOJ2243: [SDOI2011]染色（洛谷P2486）

树链剖分

BZOJ题目传送门
洛谷题目传送门

把树剖分后，查询和修改操作可以用线段树来实现。

对于线段树上的每一个节点，定义lc表示其左端点的颜色，rc表示其右端点的颜色。那么在统计颜色段数量的时候只需判断其左区间的rc是否与右区间的lc是否相等即可。如果相等则把颜色段总数-1。

在进行路径修改/查询时，要注意一些细节。比如查询时判断这条链的底和上一条链的顶的颜色。

细节决定成败。（但是BZOJ好像没有0颜色）

代码：

#include<cctype>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define N 100005
using namespace std;
struct tree{
    int l,r,lc,rc,sum,tg;
}t[N*4];
struct edge{
    int next,to;
}ed[N*2];
int n,m,nd,k,c[N],to[N],sz[N],id[N];
int in[N],h[N],dep[N],fa[N],tp[N];
inline int _read(){
    int num=0; char ch=getchar();
    while (!isdigit(ch)) ch=getchar();
    while (isdigit(ch)) num=num*10+ch-48,ch=getchar(); 
    return num;
}
inline void addedge(int x,int y){
    ed[++k].next=h[x],ed[k].to=y,h[x]=k;
}
void dfs1(int x,int depth){
    sz[x]=1,dep[x]=depth;
    for (int i=h[x];i;i=ed[i].next)
        if (ed[i].to!=fa[x]){
            int v=ed[i].to;
            fa[v]=x,dfs1(v,depth+1),sz[x]+=sz[v];
            if (sz[v]>sz[to[x]]) to[x]=v;
        }
}
void dfs2(int x){
    id[x]=++nd,in[nd]=x;
    for (int i=h[x];i;i=ed[i].next)
        if (ed[i].to==to[x])
            tp[ed[i].to]=tp[x],dfs2(ed[i].to);
    for (int i=h[x];i;i=ed[i].next)
        if (ed[i].to!=fa[x]&&ed[i].to!=to[x])
            tp[ed[i].to]=ed[i].to,dfs2(ed[i].to);
}
void pshp(tree &x,tree y,tree z){
    x.sum=y.sum+z.sum-(y.rc==z.lc);
    x.lc=y.lc,x.rc=z.rc;
}
void pshd(int x){
    t[x*2].tg=t[x*2+1].tg=t[x].tg;
    t[x*2].sum=t[x*2+1].sum=1;
    t[x*2].lc=t[x*2].rc=t[x*2+1].lc=t[x*2+1].rc=t[x].tg;
    t[x].tg=0;
}
void build(int l,int r,int x){
    t[x].l=l,t[x].r=r;
    if (l==r){ t[x].sum=1,t[x].lc=t[x].rc=c[in[l]]; return; }
    int mid=l+r>>1; build(l,mid,x*2),build(mid+1,r,x*2+1);
    pshp(t[x],t[x*2],t[x*2+1]);
}
void nsrt(int x,int l,int r,int w){
    if (t[x].l>r||t[x].r<l) return;
    if (t[x].l>=l&&t[x].r<=r){
        t[x].tg=t[x].lc=t[x].rc=w;
        t[x].sum=1; return;
    }
    if (t[x].tg) pshd(x);
    nsrt(x*2,l,r,w),nsrt(x*2+1,l,r,w);
    pshp(t[x],t[x*2],t[x*2+1]);
}
void mdfy(int x,int y,int z){
    while (tp[x]!=tp[y]){
        if (dep[tp[x]]<dep[tp[y]]) swap(x,y);
        nsrt(1,id[tp[x]],id[x],z),x=fa[tp[x]];
    }
    if (dep[x]<dep[y]) swap(x,y); nsrt(1,id[y],id[x],z);
}
int find(int x,int l,int r){
    if (t[x].l>r||t[x].r<l) return 0;
    if (t[x].l>=l&&t[x].r<=r) return t[x].sum;
    if (t[x].tg) pshd(x);
    int lx=find(x*2,l,r),rx=find(x*2+1,l,r),ans=lx+rx;
    if (lx&&rx) ans-=(t[x*2].rc==t[x*2+1].lc);
    return ans;
}
int dfs(int x,int p){//查询在区间位置上为p的颜色
    if (t[x].l==t[x].r&&t[x].l==p) return t[x].lc;
    if (t[x].tg) return t[x].tg;
    int mid=t[x].l+t[x].r>>1;
    if (p<=mid) return dfs(x*2,p); 
    return dfs(x*2+1,p);
}
int srch(int x,int y){//这个地方的细节很多
    int ans=0,p[2][2],q[2][2];//p,q分别记录x,y上一条链顶和底的颜色
    memset(p,0,sizeof(p)),memset(q,0,sizeof(q));
    while (tp[x]!=tp[y]){
        if (dep[tp[x]]<dep[tp[y]]) swap(x,y),swap(p,q);//p,q也要跟着换
        ans+=find(1,id[tp[x]],id[x]);
        p[0][0]=dfs(1,id[tp[x]]),p[0][1]=dfs(1,id[x]);//记录顶和底的颜色
        if (p[0][1]==p[1][0]) ans--; x=fa[tp[x]];
        memcpy(p[1],p[0],sizeof(p[1]));//记得更新
    }
    if (dep[x]<dep[y]) swap(x,y),swap(p,q);
    ans+=find(1,id[y],id[x]);
    if (dfs(1,id[x])==p[0][0]) ans--;//这里x和y都需要判一遍
    if (dfs(1,id[y])==q[0][0]) ans--;
    return ans;
}
int main(){
    n=_read(),m=_read();
    for (int i=1;i<=n;i++) c[i]=_read();
    for (int i=1;i<n;i++){
        int u=_read(),v=_read();
        addedge(u,v),addedge(v,u);
    }
    dfs1(1,1),tp[1]=1,dfs2(1),build(1,n,1);
    while (m--){
        char s[3]; scanf("%s",s);
        int x=_read(),y=_read(),z;
        if (s[0]=='C') z=_read(),mdfy(x,y,z);
        else printf("%d\n",srch(x,y));
    }
    return 0;
}