[树链剖分] [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]之间。

题目一看就知道是树剖，就是写起来有点烦
注意线段树询问的时候和树剖的时候都要判断边界颜色是否一样

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

const int MaxN=100005;

int N,M;
int Tot,Head[MaxN],Next[MaxN<<1],To[MaxN<<1];
int Fa[MaxN],Dep[MaxN],Size[MaxN],Son[MaxN];
int Index,Top[MaxN],ID[MaxN],Point[MaxN];
int Col[MaxN];

struct Node{
    int l,r;
    int lcol,rcol,cnt;
    int tag;
};

struct Ans{
    int cnt,lcol,rcol;
};

#define L (Seg[cur].l)
#define R (Seg[cur].r)
#define Mid (Seg[cur].l+Seg[cur].r>>1)
#define LCol (Seg[cur].lcol)
#define RCol (Seg[cur].rcol)
#define Cnt (Seg[cur].cnt)
#define Tag (Seg[cur].tag)
#define LCH (cur<<1)
#define RCH (cur<<1|1)

struct Segment_Tree{
    Node Seg[MaxN<<2];

    void PushUp(int cur){
        LCol=Seg[LCH].lcol,RCol=Seg[RCH].rcol;
        Cnt=Seg[LCH].cnt+Seg[RCH].cnt-(Seg[LCH].rcol==Seg[RCH].lcol);
    }

    void Build(int cur,int l,int r){
        L=l,R=r;
        if(L==R){
            LCol=RCol=Col[Point[l]];
            Cnt=1;
            return;
        }
        Build(LCH,L,Mid),Build(RCH,Mid+1,R);
        PushUp(cur);
    }

    void PushDown(int cur){
        if(Tag){
            Seg[LCH].lcol=Seg[LCH].rcol=Seg[RCH].lcol=Seg[RCH].rcol=Tag;
            Seg[LCH].cnt=Seg[RCH].cnt=1;
            Seg[LCH].tag=Seg[RCH].tag=Tag;
            Tag=0;
        }
    }

    void Update(int cur,int l,int r,int v){
        if(L>=l&&R<=r){
            LCol=RCol=v,Cnt=1,Tag=v;
            return;
        }
        PushDown(cur);
        if(Mid>=l)
            Update(LCH,l,r,v);
        if(Mid<r)
            Update(RCH,l,r,v);
        PushUp(cur);
    }

    int Get_Col(int cur,int pos){
        if(L==R)
            return LCol;
        PushDown(cur);
        if(Mid>=pos)
            return Get_Col(LCH,pos);
        return Get_Col(RCH,pos);
    }

    Ans Get_Cnt(int cur,int l,int r){
        if(L>=l&&R<=r)
            return (Ans){Cnt,LCol,RCol};
        PushDown(cur);
        Ans a=(Ans){0,0,0},b=(Ans){0,0,0};
        if(Mid>=l)
            a=Get_Cnt(LCH,l,r);
        if(Mid<r)
            b=Get_Cnt(RCH,l,r);
        return (Ans){a.cnt+b.cnt-(Mid>=l&&Mid<r)*(a.rcol==b.lcol),a.lcol,b.rcol};
    }
}Seg_Tree;

int Input(){
    int s=0;
    char ch;
    while(ch=getchar(),ch<'0'||ch>'9');
    while(ch>='0'&&ch<='9'){
        s=s*10+ch-'0';
        ch=getchar();
    }
    return s;
}

void Add_Edge(int u,int v){
    Next[++Tot]=Head[u];
    To[Tot]=v;
    Head[u]=Tot;
}

void DFS1(int u,int fa){
    int i,v;
    Size[u]=1;
    for(i=Head[u];i;i=Next[i])
        if((v=To[i])!=fa){
            Fa[v]=u,Dep[v]=Dep[u]+1;
            DFS1(v,u);
            Size[u]+=Size[v];
            Son[u]=(Size[Son[u]]>Size[v]?Son[u]:v);
        }
}

void DFS2(int u){
    int i,v;
    Point[ID[u]=++Index]=u;
    if(v=Son[u])
        Top[v]=Top[u],DFS2(v);
    for(i=Head[u];i;i=Next[i])
        if((v=To[i])!=Fa[u]&&v!=Son[u])
            Top[v]=v,DFS2(v);
}

void Update(int u,int v,int c){
    int r1=Top[u],r2=Top[v];
    while(r1!=r2){
        if(Dep[r1]<Dep[r2])
            swap(u,v),swap(r1,r2);
        Seg_Tree.Update(1,ID[r1],ID[u],c);
        u=Fa[r1],r1=Top[u];
    }
    if(Dep[u]>Dep[v])   
        swap(u,v);
    Seg_Tree.Update(1,ID[u],ID[v],c);
}

void Query(int u,int v){
    int r1=Top[u],r2=Top[v],cnt=0;
    while(r1!=r2){
        if(Dep[r1]<Dep[r2])
            swap(u,v),swap(r1,r2);
        cnt+=Seg_Tree.Get_Cnt(1,ID[r1],ID[u]).cnt;
        cnt-=(Fa[r1]&&Seg_Tree.Get_Col(1,ID[r1])==Seg_Tree.Get_Col(1,ID[Fa[r1]]));
        u=Fa[r1],r1=Top[u];
    }
    if(Dep[u]>Dep[v])
        swap(u,v);
    cnt+=Seg_Tree.Get_Cnt(1,ID[u],ID[v]).cnt;
    printf("%d\n",cnt);
}

int main(){
    int i,u,v,c;
    char opt[5];
    N=Input(),M=Input();
    for(i=1;i<=N;i++)
        Col[i]=Input();
    for(i=1;i<N;i++){
        u=Input(),v=Input();
        Add_Edge(u,v);
        Add_Edge(v,u);
    }
    DFS1(1,0);
    Top[1]=1,DFS2(1);
    Seg_Tree.Build(1,1,N);
    while(M--){
        scanf("%s",opt),u=Input(),v=Input();
        if(*opt=='C'){
            c=Input();
            Update(u,v,c);
        }
        else Query(u,v);
    }
    return 0;
}