[BZOJ]2243 染色 树链剖分+线段树

题目链接：http://www.lydsy.com/JudgeOnline/problem.php?id=2243
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]之间。

这道题是一道树链剖分的好题，感觉对树链剖分理解的更透彻了，并且能用于线段树中.
修改，统计树上的一条链，自然想到树链剖分，这里用到的是重链剖分.我们dfs两次求出时间戳之后（本题只有链，所以out数组无所谓），用线段树进行维护，要维护的信息有本区间最左边的颜色lc，最右边的颜色rc，颜色段数量c,c这个要把左儿子的c加上右儿子的c，并判断左儿子的rc和右儿子的lc是否一样，若一样的话本区间要c–；
在查询的时候我们要注意除了ans加上本条链的c，还要判断本条链最下面的颜色与上次这条链所在路径处理的上一条链的最上面的颜色作比较，若一样还要ans–（因为统计的是颜色段）.在u和top一样的时候还是要比较.
备注：
1. query线段树的时候若两儿子都有值则要判断相接地方颜色是否一样，若一样还是要–.
详见代码.

/**************************************************************
    Problem: 2243
    User: MaxMercer
    Language: C++
    Result: Accepted
    Time:4536 ms
    Memory:35648 kb
****************************************************************/

#include<stdio.h>
#include<algorithm>
using namespace std;
const int maxn=200005;
inline const int read(){
    register int f=1,x=0;
    register char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=(x<<3)+(x<<1)+ch-'0';ch=getchar();}
    return x*f;
}
int h[maxn],dep[maxn],top[maxn],siz[maxn],fa[maxn],indexx,seq[maxn],in[maxn],out[maxn],son[maxn],a[maxn],num,n,T,Rc,Lc;
char ss[5];
struct node{
    node *ls,*rs;
    int lc,rc,c,flag;
    void pushdown(){
         if(flag!=-1){
            ls->lc=flag,ls->rc=flag;
            ls->c=1,ls->flag=flag;
            rs->lc=flag,rs->rc=flag;
            rs->c=1,rs->flag=flag;
            flag=-1;
         }
    }
    void update(){
         if(ls->rc!=rs->lc) c=ls->c+rs->c;
         else c=ls->c+rs->c-1;
         lc=ls->lc,rc=rs->rc;
    }
}pool[4*maxn],*tail=pool,*root;
struct edge{
    int v,nxt;
}e[3*maxn]; 
void add(int u,int v){
    e[++num].v=v;
    e[num].nxt=h[u];
    h[u]=num;
}
void dfs1(int u){
    siz[u]=1;
    for(int i=h[u];i;i=e[i].nxt){
        int v=e[i].v;
        if(v==fa[u]) continue;
        fa[v]=u;
        dep[v]=dep[u]+1;
        dfs1(v);
        siz[u]+=siz[v];
        if(siz[son[u]]<siz[v]) son[u]=v;
    }
}
void dfs2(int u,int tp){
    top[u]=tp;
    seq[++indexx]=u;
    in[u]=indexx;
    if(son[u]) dfs2(son[u],tp);
    for(int i=h[u];i;i=e[i].nxt){
        int v=e[i].v;
        if(v==fa[u]||v==son[u]) continue;
        dfs2(v,v);
    }
    out[u]=indexx;
}
node *build(int lf,int rg){
    node *bt=++tail;
    if(lf==rg) {bt->lc=bt->rc=a[seq[lf]];bt->c=1;bt->flag=-1;return bt;}
    int mid=(lf+rg)>>1;
    bt->ls=build(lf,mid);
    bt->rs=build(mid+1,rg);
    bt->c=bt->ls->c+bt->rs->c;
    if(bt->ls->rc==bt->rs->lc) bt->c--;
    bt->lc=bt->ls->lc;
    bt->rc=bt->rs->rc;
    bt->flag=-1;
    return bt;
}
void modify(node *bt,int lf,int rg,int L,int R,int delta){
    if(L<=lf&&rg<=R) {bt->c=1;bt->lc=delta;bt->rc=delta;bt->flag=delta;return;}
    int mid=(lf+rg)>>1;
    bt->pushdown();
    if(L<=mid) modify(bt->ls,lf,mid,L,R,delta);
    if(R>mid) modify(bt->rs,mid+1,rg,L,R,delta);
    bt->update();
    return; 
}
int query(node *bt,int lf,int rg,int L,int R){
    int bin;
    if(L==lf) Lc=bt->lc;
    if(R==rg) Rc=bt->rc;
    if(L<=lf&&rg<=R) return bt->c;
    int mid=(lf+rg)>>1;
    bt->pushdown();
    int rt=0;
    if(L<=mid) rt+=query(bt->ls,lf,mid,L,R);
    bin=rt;
    if(R>mid) rt+=query(bt->rs,mid+1,rg,L,R);
    if(rt>bin&&bin&&bt->ls->rc==bt->rs->lc) rt--;  
    bt->update();
    return rt; 
}
inline void qu_modify(int u,int v,int val){
    while(top[u]!=top[v]){
        if(dep[top[u]]<dep[top[v]]) swap(u,v);
        modify(root,1,n,in[top[u]],in[u],val);
        u=fa[top[u]];
    }
    if(dep[u]<dep[v]) swap(u,v);
    modify(root,1,n,in[v],in[u],val);
}
inline int qu_query(int u,int v){
    int rt=0,ans1=-1,ans2=-1;
    Rc=-2,Lc=-2;
    while(top[u]!=top[v]){
       if(dep[top[u]]<dep[top[v]]) {swap(u,v);swap(ans1,ans2);} 
       rt+=query(root,1,n,in[top[u]],in[u]);
       if(Rc==ans1) rt--;
       ans1=Lc;
       u=fa[top[u]];
    }
    if(dep[u]<dep[v]) {swap(u,v);swap(ans1,ans2);}
    rt+=query(root,1,n,in[v],in[u]);
    if(Rc==ans1) rt--;
    if(Lc==ans2) rt--;
    return rt;
}
int main(){
    int x,y;
    n=read(),T=read();
    for(register int i=1;i<=n;i++) a[i]=read(); 
    for(register int i=1;i<n;i++) x=read(),y=read(),add(x,y),add(y,x);
    dep[1]=fa[1]=1;
    dfs1(1);
    dfs2(1,1);
    root=build(1,n);
    while(T--){
       int u,v,delta;
       scanf("%s",ss);
       if(ss[0]=='C'){
          u=read(),v=read(),delta=read();
          qu_modify(u,v,delta);
       }
       if(ss[0]=='Q'){
          u=read(),v=read();
          printf("%d\n",qu_query(u,v));
       }
    }
}