P2486 [SDOI2011]染色 （树链剖分）

P2486 [SDOI2011]染色

分析：

颜色段数用线段树维护，记录最左端颜色与最右端颜色，每一次合并的时候，如果中间颜色有重叠，就将段数--。

链剖查询的时候注意两条链端点颜色相同的情况。

因为跳链是一个交替的过程，记录一下左右链跳完之后跳到的颜色是哪一种：ans1对应左边的x，ans2对应右边的y。

每一次x与y深度交换的时候，将ans1与ans2也交换，保证对应的关系。

#include<bits/stdc++.h>
using namespace std;
#define ri register int 
#define N 100005
#define mid ((l+r)>>1)
int Lc,Rc,lc[N*4],rc[N*4],sum[N*4],add[N*4],n;
int fa[N],siz[N],son[N],id[N],dfn[N],cnt=0,dep[N],col[N],topp[N];
vector<int> e[N];
void dfs1(int u,int ff)
{
    siz[u]=1;
    for(ri i=0;i<e[u].size();++i){
        int v=e[u][i];
        if(v==ff) continue; fa[v]=u;
        dep[v]=dep[u]+1; dfs1(v,u); siz[u]+=siz[v];
        if(siz[son[u]]<siz[v]) son[u]=v;
    }
}
void dfs2(int u,int tp)
{
    dfn[++cnt]=u; id[u]=cnt; topp[u]=tp;
    if(son[u]) dfs2(son[u],tp);
    for(ri i=0;i<e[u].size();++i){
        int v=e[u][i];
        if(v==fa[u] || v==son[u]) continue;
        dfs2(v,v);
    }
}
void update(int s)
{
    sum[s]=sum[s<<1]+sum[s<<1|1]-(rc[s<<1]==lc[s<<1|1]);
    lc[s]=lc[s<<1]; rc[s]=rc[s<<1|1];
}
void build(int s,int l,int r)
{
    if(l==r) { sum[s]=1; add[s]=0; lc[s]=rc[s]=col[dfn[l]]; return ; }
    build(s<<1,l,mid); build(s<<1|1,mid+1,r);
    update(s);
}
void pushdown(int s)
{
    if(!add[s]) return ;
    int v=add[s];
    add[s<<1]=add[s<<1|1]=lc[s<<1]=lc[s<<1|1]=rc[s<<1]=rc[s<<1|1]=v;
    sum[s<<1]=sum[s<<1|1]=1;
    add[s]=0;
}
void modify(int s,int l,int r,int L,int R,int v)
{
    if(L<=l && r<=R) { add[s]=v; sum[s]=1; lc[s]=rc[s]=v; return;  }
    pushdown(s);
    if(L<=mid) modify(s<<1,l,mid,L,R,v);
    if(R>mid)  modify(s<<1|1,mid+1,r,L,R,v);
    update(s);
}
int query(int s,int l,int r,int L,int R)
{
    if(l==L) Lc=lc[s]; if(r==R) Rc=rc[s];
    if(L<=l && r<=R) return sum[s];
    pushdown(s);
    int ans=0,fl=0;
    if(L<=mid) ans+=query(s<<1,l,mid,L,R),fl++;
    if(R>mid)  ans+=query(s<<1|1,mid+1,r,L,R),fl++;
    if(fl==2) ans-=(rc[s<<1]==lc[s<<1|1]);
    return ans;
}
void modify_link(int x,int y,int c)
{
    while(topp[x]!=topp[y]){
        if(dep[topp[x]]<dep[topp[y]]) swap(x,y);
        modify(1,1,n,id[topp[x]],id[x],c);
        x=fa[topp[x]];
    }
    if(dep[x]<dep[y]) swap(x,y);
    modify(1,1,n,id[y],id[x],c);
}
int query_link(int x,int y)
{
    int ans=0,ans1=-1,ans2=-1;
    while(topp[x]!=topp[y]){
        if(dep[topp[x]]<dep[topp[y]]) swap(x,y),swap(ans1,ans2);//x对应ans1 y对应ans2 
        ans+=query(1,1,n,id[topp[x]],id[x]);
        if(ans1==Rc) ans--;//Rc记录的是这一次链的下端 如果和上一次的上端颜色是一样的 就-- 
        ans1=Lc; x=fa[topp[x]];
    }
    if(dep[x]>dep[y]) swap(x,y),swap(ans1,ans2);//!!ans1一定是和x对应的 所以说如果下面要用ans1与Lc比较的话 要保证x在上方 
    ans+=query(1,1,n,id[x],id[y]);
    if(ans1==Lc) ans--;
    if(ans2==Rc) ans--;
    return ans;
}
int main()
{
    int a,b,c,m;
    scanf("%d%d",&n,&m);
    for(ri i=1;i<=n;++i) scanf("%d",&col[i]);
    for(ri i=1;i<=n-1;++i) scanf("%d%d",&a,&b),e[a].push_back(b),e[b].push_back(a);
    dep[1]=1; dfs1(1,0); dfs2(1,1);
    build(1,1,n);
    char op[2]; 
    for(ri i=1;i<=m;++i){
        scanf("%s%d%d",op,&a,&b);
        if(op[0]=='C') scanf("%d",&c),modify_link(a,b,c);
        else printf("%d\n",query_link(a,b));
    }
    return 0;
}
/*
6 100
2 2 1 2 1 1
1 2
1 3
2 4
2 5
2 6
Q 3 4
Q 1 5
Q 6 5
Q 6 4
C 1 4 1
Q 6 4


Q 3 5
C 2 1 1
Q 3 5
C 5 1 2
Q 3 5
*/

View Code