【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]之间。

题解：
树链剖分+线段树。比较好的一道题，主要是细节要注意，也就是接驳处颜色是否相同。树剖后，线段树要记录左端点l,右端点r，左端点的颜色lc,右端点的颜色rc,区间成段更新的标记tag,区间有多少颜色段。区间合并的时候要注意如果左子树的右端和右子树的左端颜色相同那么数量要减一。但是存在一个问题当前剖到的链与上一次的链在相交的边缘可能颜色相同，如果颜色相同答案需要减一。所以统计答案的时候要记录下上一次剖到的链的左端点的颜色，与当前剖到的链右端点的颜色（因为在处理出的线段树中越靠近根的点位置越左），比较这两个颜色，若相同则答案减1。

代码如下：

#include<iostream>
#include<stdio.h>
#include<algorithm>
#include<string.h>
#define N 100005
#define ll long long
#define inf 0x7f7f7f7f
#define ls (k<<1)
#define rs (k<<1|1)
#define mid (l+r>>1)
using namespace std;
int n,m,cnt,sz,hd[N],dep[N],siz[N],bl[N],pos[N],v[N],fa[N][18];
struct seg
{
    int l,r,lc,rc,s,tag;
}t[N*3];
struct edge
{
    int to,nex;
}e[N<<1];
void insert(int u,int v)
{
    e[++cnt].to=v,e[cnt].nex=hd[u],hd[u]=cnt;
    e[++cnt].to=u,e[cnt].nex=hd[v],hd[v]=cnt;
}
void dfs1(int x,int f)
{
    siz[x]=1;
    for(int i=1;i<=17 && (1<<i)<=dep[x];i++)
    fa[x][i]=fa[fa[x][i-1]][i-1];
    for(int i=hd[x];i;i=e[i].nex)
    {
        if(e[i].to==f) continue;
        dep[e[i].to]=dep[x]+1;
        fa[e[i].to][0]=x;
        dfs1(e[i].to,x);
        siz[x]+=siz[e[i].to];
    }
}
void dfs2(int x,int top)
{
    pos[x]=++sz;bl[x]=top;
    int k=0;
    for(int i=hd[x];i;i=e[i].nex)
    if(dep[e[i].to]>dep[x] && siz[k]<siz[e[i].to]) k=e[i].to;
    if(k) dfs2(k,top);
    for(int i=hd[x];i;i=e[i].nex)
    if(dep[e[i].to]>dep[x] && k!=e[i].to) dfs2(e[i].to,e[i].to);
}
int lca(int x,int y)
{
    if(dep[x]<dep[y]) swap(x,y);
    int t=dep[x]-dep[y];
    for(int i=0;i<=17;i++)
    if(t&(1<<i)) x=fa[x][i];
    for(int i=17;i>=0;i--)
    if(fa[x][i]!=fa[y][i])
    x=fa[x][i],y=fa[y][i];
    if(x==y) return x;
    return fa[x][0];
}
void build(int k,int l,int r)
{
    t[k].l=l;t[k].r=r;t[k].s=1;t[k].tag=-1;
    if(l==r) return;
    build(ls,l,mid);build(rs,mid+1,r);
}
void pushup(int k)
{
    t[k].lc=t[ls].lc;t[k].rc=t[rs].rc;
    if(t[ls].rc^t[rs].lc) t[k].s=t[ls].s+t[rs].s;
    else t[k].s=t[ls].s+t[rs].s-1;
}
void pushdown(int k)
{
    int tmp=t[k].tag;t[k].tag=-1;
    if(tmp==-1 || t[k].l==t[k].r)return;
    t[ls].s=t[rs].s=1;
    t[ls].tag=t[rs].tag=tmp;
    t[ls].lc=t[ls].rc=tmp;
    t[rs].lc=t[rs].rc=tmp;
}
void change(int k,int L,int R,int c)
{
    pushdown(k);
    int l=t[k].l,r=t[k].r;
    if(l==L && r==R)
    {
        t[k].lc=t[k].rc=c;
        t[k].s=1;t[k].tag=c;
        return;
    }
    if(mid>=R) change(ls,L,R,c);
    else if(mid<L) change(rs,L,R,c);
    else change(ls,L,mid,c),change(rs,mid+1,R,c);
    pushup(k);
}
int ask(int k,int L,int R)
{
    pushdown(k);
    int l=t[k].l,r=t[k].r;
    if(l==L && r==R) return t[k].s;
    if(mid>=R) return ask(ls,L,R);
    else if(mid<L) return ask(rs,L,R);
    else
    {
        int tmp=1;
        if(t[ls].rc^t[rs].lc) tmp=0;
        return ask(ls,L,mid)+ask(rs,mid+1,R)-tmp;
    }
}
int getcol(int k,int x)
{
    pushdown(k);
    int l=t[k].l,r=t[k].r;
    if(l==r) return t[k].lc;
    if(x<=mid) return getcol(ls,x);
    else return getcol(rs,x);
}
int solvesum(int x,int f)
{
    int sum=0;
    while(bl[x]!=bl[f])
    {
        sum+=ask(1,pos[bl[x]],pos[x]);
        if(getcol(1,pos[bl[x]])==getcol(1,pos[fa[bl[x]][0]])) sum--;
        x=fa[bl[x]][0]; 
    }
    sum+=ask(1,pos[f],pos[x]);
    return sum;
}
void solvechange(int x,int f,int c)
{
    while(bl[x]!=bl[f])
    {
        change(1,pos[bl[x]],pos[x],c);
        x=fa[bl[x]][0]; 
    }
    change(1,pos[f],pos[x],c);
}
int main()
{
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++) scanf("%d",&v[i]);
    for(int i=1;i<n;i++)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        insert(x,y);
    }
    int a,b,c;
    dfs1(1,0);dfs2(1,1);
    build(1,1,n);
    for(int i=1;i<=n;i++) change(1,pos[i],pos[i],v[i]);
    for(int i=1;i<=m;i++)
    {
        char ch[10];
        scanf("%s",ch);
        if(ch[0]=='Q')
        {
            scanf("%d%d",&a,&b);
            int t=lca(a,b);
            printf("%d\n",solvesum(a,t)+solvesum(b,t)-1);
        }
        else
        {
            scanf("%d%d%d",&a,&b,&c);
            int t=lca(a,b);
            solvechange(a,t,c);solvechange(b,t,c);
        }
    }
    return 0;
}