[bzoj2243][SDOI2011]染色

Description

给定一棵有\(n\)个节点的无根树和\(m\)个操作,操作有\(2\)类:

1.将节点\(a\)到节点\(b\)路径上所有点都染成颜色\(c\);

2.询问节点\(a\)到节点\(b\)路径上的颜色段数量(连续相同颜色被认为是同一段),如\("112221"\)由\(3\)段组成:\("11","222"和"1"\).

请你写一个程序依次完成这\(m\)个操作.

Input

第一行包含\(2\)个整数\(n\)和\(m\),分别表示节点数和操作数.

第二行包含\(n\)个正整数表示\(n\)个节点的初始颜色.

下面\(n-1\)行每行包含两个整数\(x\)和\(y\),表示\(x\)和\(y\)之间有一条无向边.

下面\(m\)行每行描述一个操作:

\(”C\;a\;b\;c”\)表示这是一个染色操作,把节点\(a\)到节点\(b\)路径上所有点(包括\(a\)和\(b\))都染成颜色\(c\);

\(”Q\;a\;b”\)表示这是一个询问操作,询问节点\(a\)到节点\(b\)(包括\(a\)和\(b\))路径上的颜色段数量.

Output

对于每个询问操作,输出一行答案.

Sample Input

6 5 
2 2 1 2 1 1 
1 2 
1 3 
2 4 
2 5 
2 6 
Q 3 5 
C 2 1 1 
Q 3 5 
C 5 1 2 
Q 3 5

Sample Output

3 
1 
2

HINT

\(N,M\;\leq\;10^5,c\;\in\;z,c\;\in\;[0, 10^9]\).

Solution

树链剖分+线段树.

每次上移的时候判断相邻两段端点是否同色即可.

#include<cmath>
#include<ctime>
#include<queue>
#include<stack>
#include<cstdio>
#include<vector>
#include<cstring>
#include<cstdlib>
#include<iostream>
#include<algorithm>
#define N 100005
#define M 300005
using namespace std;
struct linetree{
    int l,r,lc,rc,cnt;
}lt[M];
struct graph{
    int nxt,to;
}e[M];
char c[2];
int g[N],a[N],w[N],n,m,u,v,k,l,cnt;
int f[N],p[N],dep[N],siz[N],son[N],top[N];
inline void addedge(int x,int y){
    e[++cnt].nxt=g[x];g[x]=cnt;e[cnt].to=y;
}
inline void dfs1(int u){
    int m=0;siz[u]=1;
    for(int i=g[u];i;i=e[i].nxt)
        if(!dep[e[i].to]){
            f[e[i].to]=u;
            dep[e[i].to]=dep[u]+1;
            dfs1(e[i].to);
            siz[u]+=siz[e[i].to];
            if(siz[e[i].to]>m){
                son[u]=e[i].to;
                m=siz[e[i].to];
            } 
        }
}
inline void dfs2(int u,int tp){
    top[u]=tp;p[u]=++cnt;w[cnt]=a[u];
    if(son[u]) dfs2(son[u],tp);
    for(int i=g[u];i;i=e[i].nxt)
        if(f[u]!=e[i].to&&e[i].to!=son[u])
            dfs2(e[i].to,e[i].to);
}
inline void build(int u,int l,int r){
    lt[u].l=l;lt[u].r=r;
    if(lt[u].l<lt[u].r){
        int lef=u<<1,rig=u<<1|1;
        int mid=(lt[u].l+lt[u].r)>>1;
        build(lef,l,mid);build(rig,mid+1,r);
        lt[u].cnt=lt[lef].cnt+lt[rig].cnt;
        if(lt[lef].rc==lt[rig].lc) --lt[u].cnt;
        lt[u].lc=lt[lef].lc;lt[u].rc=lt[rig].rc;
    }
    else{
        lt[u].cnt=1;lt[u].lc=lt[u].rc=w[lt[u].l];
    }
}
inline void cover(int u,int l,int r,int k){
    if(lt[u].l>=l&&lt[u].r<=r){
        lt[u].lc=lt[u].rc=k;lt[u].cnt=1;
    }
    else if(lt[u].l<lt[u].r){
        int lef=u<<1,rig=u<<1|1;
        int mid=(lt[u].l+lt[u].r)>>1;
        if(lt[u].cnt==1){
            lt[lef].cnt=lt[rig].cnt=1;
            lt[lef].lc=lt[lef].rc=lt[rig].lc=lt[rig].rc=lt[u].lc;
        }
        if(l<=mid) cover(lef,l,r,k);
        if(r>mid) cover(rig,l,r,k);
        lt[u].cnt=lt[lef].cnt+lt[rig].cnt;
        if(lt[lef].rc==lt[rig].lc) --lt[u].cnt;
        lt[u].lc=lt[lef].lc;lt[u].rc=lt[rig].rc;
    }
}
inline int ask(int u,int l,int r){
    if(lt[u].l>=l&&lt[u].r<=r)
        return lt[u].cnt;
    if(lt[u].l<lt[u].r){
        int lef=u<<1,rig=u<<1|1,ret=0;
        int mid=(lt[u].l+lt[u].r)>>1;
        if(lt[u].cnt==1){
            lt[lef].cnt=lt[rig].cnt=1;
            lt[lef].lc=lt[rig].lc=lt[lef].rc=lt[rig].rc=lt[u].lc;
        }
        if(l<=mid) ret+=ask(lef,l,r);
        if(r>mid) ret+=ask(rig,l,r);
        if(l<=mid&&r>mid&&lt[lef].rc==lt[rig].lc) --ret;
        return ret;
    }
}
inline int col(int u,int x){
    if(lt[u].l==x) return lt[u].lc;
    if(lt[u].r==x) return lt[u].rc;
    if(lt[u].l<lt[u].r){
        int lef=u<<1,rig=u<<1|1;
        int mid=(lt[u].l+lt[u].r)>>1;
        if(lt[u].cnt==1){
            lt[lef].cnt=lt[rig].cnt=1;
            lt[lef].lc=lt[rig].lc=lt[lef].rc=lt[rig].rc=lt[u].lc;
        }
        if(x<=mid) return col(lef,x);
        return col(rig,x);
    }
}
inline void cov(int x,int y,int k){
    int t;
    while(top[x]!=top[y]){    
        if(dep[top[x]]<dep[top[y]]){
            t=x;x=y;y=t;
        }
        cover(1,p[top[x]],p[x],k);
        x=f[top[x]]; 
    }
    if(p[x]>p[y]){
        t=x;x=y;y=t;
    }
    cover(1,p[x],p[y],k);
}
inline int q(int x,int y){
    int ret=0,t;
    while(top[x]!=top[y]){
        if(dep[top[x]]<dep[top[y]]){
            t=x;x=y;y=t;
        }
        ret+=ask(1,p[top[x]],p[x]);
        if(col(1,p[top[x]])==col(1,p[f[top[x]]])) --ret;
        x=f[top[x]];
    }
    if(p[x]>p[y]){
        t=x;x=y;y=t;
    }
    ret+=ask(1,p[x],p[y]);
    return ret;
}
inline int lca(int x,int y){
    int t;
    while(top[x]!=top[y]){    
        if(dep[top[x]]<dep[top[y]]){
            t=x;x=y;y=t;
        }
        x=f[top[x]]; 
    }
    if(p[x]>p[y]){
        t=x;x=y;y=t;
    }
    return x;
}
inline void Aireen(){
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;++i)
        scanf("%d",&a[i]);
    for(int i=1,x,y;i<n;++i){
        scanf("%d%d",&x,&y);
        addedge(x,y);addedge(y,x);
    }
    dep[1]=1;dfs1(1);
    cnt=0;dfs2(1,1);
    build(1,1,n);
    while(m--){
        scanf("%s%d%d",c,&u,&v);
        if(c[0]=='C'){
            scanf("%d",&k);cov(u,v,k); 
        }
        else{
            l=lca(u,v);
            printf("%d\n",q(u,v));
        }
    }
}
int main(){
    freopen("color.in","r",stdin);
    freopen("color.out","w",stdout);
    Aireen();
    fclose(stdin);
    fclose(stdout);
    return 0;
}