[题解]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个节点的初始颜色
下面 行每行包含两个整数x和y，表示x和y之间有一条无向边。
下面 行每行描述一个操作：
“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<=10^5，操作数M<=10^5，所有的颜色C为整数且在[0, 10^9]之间。

Solution

树链剖分即可，颜色段数还是很好用线段树在区间维护的，就是注意一下树链合并的时候是有方向性的，最后合并的时候要把其中一条链反过来（自己脑补一下）。

代码：

#include<cstdio>
#include<algorithm>
using namespace std;

template<typename T>inline void read(T &x){
    T f=1;char ch=getchar();
    for(;ch<'0'||ch>'9';ch=getchar())if(ch=='-')f=-1;
    for(x=0;ch>='0'&&ch<='9';ch=getchar())x=x*10+ch-'0';
    x*=f;
}

const int maxn=100010;
struct edge{
    int to,nxt;
}e[maxn<<1];
struct Data{
    int cl,cr,num;
    Data(){cl=cr=num=0;}
    friend Data operator+(Data a,Data b){
        if((!a.num)||(!b.num))return (!a.num)?b:a;
        Data x;
        x.cl=a.cl;x.cr=b.cr;
        x.num=a.num+b.num-(a.cr==b.cl);
        return x;
    }
};
struct Seg_Tree{
    #define lc x<<1
    #define rc x<<1|1
    int L[maxn<<2],R[maxn<<2],same[maxn<<2];
    Data T[maxn<<2];
    void Build(int x,int *a,int l,int r){
        same[x]=-1;
        if((L[x]=l)==(R[x]=r)){
            T[x].cl=T[x].cr=a[l];
            return T[x].num=1,void();
        }
        int mid=(l+r)>>1;
        Build(lc,a,l,mid);Build(rc,a,mid+1,r);
        T[x]=T[lc]+T[rc];
    }
    void pushsame(int x,int c){
        T[x].cl=T[x].cr=same[x]=c;
        T[x].num=1;
    }
    void pushdown(int x){
        if(same[x]==-1)return;
        pushsame(lc,same[x]);
        pushsame(rc,same[x]);
        same[x]=-1;
    }
    void Modify(int x,int l,int r,int c){
        if(L[x]>=l&&R[x]<=r)return pushsame(x,c),void();
        int mid=(L[x]+R[x])>>1;
        pushdown(x);
        if(l<=mid)Modify(lc,l,r,c);
        if(r>mid)Modify(rc,l,r,c);
        T[x]=T[lc]+T[rc];
    }
    Data Query(int x,int l,int r){
        if(L[x]>=l&&R[x]<=r)return T[x];
        int mid=(L[x]+R[x])>>1;
        pushdown(x);
        if(l<=mid&&r>mid)return Query(lc,l,r)+Query(rc,l,r);
        else if(l<=mid)return Query(lc,l,r);
        else return Query(rc,l,r);
    }
}tree;
int n,m,num,head[maxn],col[maxn],dep[maxn];
int fa[maxn],son[maxn],top[maxn],size[maxn],rnk[maxn],pos[maxn],tot;

void add(int u,int v){
    e[++num].to=v;e[num].nxt=head[u];head[u]=num;
    e[++num].to=u;e[num].nxt=head[v];head[v]=num;
}
void Dfs(int x){
    size[x]=1;dep[x]=dep[fa[x]]+1;
    for(int i=head[x];i;i=e[i].nxt)
        if(e[i].to!=fa[x]){
            fa[e[i].to]=x;
            Dfs(e[i].to);
            size[x]+=size[e[i].to];
            if(size[e[i].to]>size[son[x]])
                son[x]=e[i].to;
        }
}
void Dfs(int x,int tp){
    top[x]=tp;
    pos[rnk[x]=++tot]=col[x];
    if(son[x])Dfs(son[x],tp);
    for(int i=head[x];i;i=e[i].nxt)
        if(e[i].to!=fa[x]&&e[i].to!=son[x])
            Dfs(e[i].to,e[i].to);
}
void Modify(int u,int v,int c){
    int x=top[u],y=top[v];
    while(x!=y){
        if(dep[x]>dep[y]){
            tree.Modify(1,rnk[x],rnk[u],c);
            x=top[u=fa[x]];
        }
        else{
            tree.Modify(1,rnk[y],rnk[v],c);
            y=top[v=fa[y]];
        }
    }
    if(dep[u]<dep[v])tree.Modify(1,rnk[u],rnk[v],c);
    else tree.Modify(1,rnk[v],rnk[u],c);
}
int Query(int u,int v){
    int x=top[u],y=top[v];
    Data ur,vr,ans;
    while(x!=y){
        if(dep[x]>dep[y]){
            ur=tree.Query(1,rnk[x],rnk[u])+ur;
            x=top[u=fa[x]];
        }
        else{
            vr=tree.Query(1,rnk[y],rnk[v])+vr;
            y=top[v=fa[y]];
        }
    }
    if(dep[u]<dep[v])vr=tree.Query(1,rnk[u],rnk[v])+vr;
    else ur=tree.Query(1,rnk[v],rnk[u])+ur;
    swap(ur.cl,ur.cr);ans=ur+vr;
    return ans.num;
}

int main(){
    read(n);read(m);
    for(int i=1;i<=n;i++)read(col[i]);
    for(int i=1,u,v;i<n;i++){
        read(u);read(v);
        add(u,v);
    }
    Dfs(1);Dfs(1,1);
    tree.Build(1,pos,1,tot);
    while(m--){
        char opt[5];
        int u,v,c;
        scanf("%s",opt);
        read(u);read(v);
        if(opt[0]=='C')read(c),Modify(u,v,c);
        else printf("%d\n",Query(u,v));
    }
    return 0;
}