【bzoj2243】【树链剖分】【线段树】SDOI2011染色

• 题目大意：
  给定一棵有N 个节点的无根树和 M 个操作，操作有2类：
  
  1. 将节点A 到节点B路径上所有点都染成颜色 C
  2. 询问节点A到节点B路径上的颜色段数量（连续相同颜色被认为是同一段），如”112221”由3段组成:”11”.”222”和”1”。
     请你写一个程序依次完成这 M 个操作。
• 数据范围：
  

• 题解：
  对树上的路径进行操作， 区间覆盖，区间查询，一瞅就可以树剖，剖完之后用线段树维护区间某种颜色出现的个数col，区间左端点颜色lcol，区间右端点颜色rcol
• 合并时满足：
  $$p->col=p->lch->col+p->rch->col-(p->lch->rcol==p->rch->lcol)$$
• 注意延重链向上爬的时候也要维护lcol和rcol进行统计答案
  直接上代码

#include <iostream>
#include <cstring>
#include <cstdio>
#include <vector>
#define MAXN 100001
using namespace std;
vector<int> e[MAXN];
int n,m;
int dep[MAXN],size[MAXN],hson[MAXN],fa[MAXN];
void dfs1(int u,int f,int d){
    size[u]=1,fa[u]=f,dep[u]=d;
    for(vector<int>::iterator it=e[u].begin();it!=e[u].end();it++){
        if(*it==f) continue;
        dfs1(*it,u,d+1);
        size[u]+=size[*it];
        if(size[*it]>size[hson[u]]) hson[u]=*it;
    }
}
int cnt[MAXN],top[MAXN],T,bel[MAXN];
void dfs2(int u,int tp){
    bel[T]=u,cnt[u]=T++,top[u]=tp;
    if(hson[u]) dfs2(hson[u],tp);
    for(vector<int>::iterator it=e[u].begin();it!=e[u].end();it++){
        if(*it==hson[u]||*it==fa[u]) continue;
        dfs2(*it,*it);
    }
}
struct node{
    int lnum,rnum,kind,l,r,mid,lzy;
    node *lch,*rch;
    node(){}
    node(int a,int b,int c):l(a),r(b),mid(c){lnum=rnum=kind=lzy=0;lch=rch=NULL;}
}*root;
int C[MAXN];
void push_down(node *p){
    p->lch->lzy=p->rch->lzy=p->lzy;
    p->lch->lnum=p->lch->rnum=p->lzy;
    p->lch->kind=1;
    p->rch->lnum=p->rch->rnum=p->lzy;
    p->rch->kind=1;
    p->lzy=0;
}
void push_up(node *p){
    p->kind=p->lch->kind+p->rch->kind;
    if(p->lch->rnum==p->rch->lnum)
        p->kind--;
    p->lnum=p->lch->lnum;
    p->rnum=p->rch->rnum;
}
node *build(int left,int right){
    node *p=new node(left,right,(left+right)>>1);
    if(right-left==1){
        p->kind=1;
        p->lnum=p->rnum=C[bel[left]];
        return p;
    }
    p->lch=build(left,p->mid);
    p->rch=build(p->mid,right);
    push_up(p);
    return p;
}
void change(node *p,int left,int right,int col){
    if(p->l==left&&p->r==right){
        p->kind=1,p->lnum=p->rnum=col;
        p->lzy=col;
        return;
    }
    if(p->lzy)
        push_down(p);
    if(right<=p->mid)
        change(p->lch,left,right,col);
    else if(left>=p->mid)
        change(p->rch,left,right,col);
    else{
        change(p->lch,left,p->mid,col);
        change(p->rch,p->mid,right,col);
    }
    push_up(p);
}
struct Q{
    int sum,lcol,rcol;
    Q(int a,int b,int c):sum(a),lcol(b),rcol(c){}
    Q(){}
};
Q query(node *p,int left,int right){
    if(p->l==left&&p->r==right)
        return Q(p->kind,p->lnum,p->rnum);
    if(p->lzy)
        push_down(p);
    if(right<=p->mid)
        return query(p->lch,left,right);
    else if(left>=p->mid)
        return query(p->rch,left,right);
    else{
        Q a=query(p->lch,left,p->mid);
        Q b=query(p->rch,p->mid,right);
        if(a.rcol==b.lcol)
            return Q(a.sum+b.sum-1,a.lcol,b.rcol);
        else
            return Q(a.sum+b.sum,a.lcol,b.rcol);
    }
}
void changeans(int u,int v,int c){
    while(top[u]!=top[v]){
        if(dep[top[u]]<dep[top[v]]) swap(u,v);
        change(root,cnt[top[u]],cnt[u]+1,c);
        u=fa[top[u]];
    }
    if(cnt[u]>cnt[v]) swap(u,v);
    change(root,cnt[u],cnt[v]+1,c);
}
int getans(int s,int t){
    pair<int,int> u,v;
    u=make_pair(s,0);
    v=make_pair(t,1);
    int lc[2]={-1};
    int ans[2]={0};
    while(top[u.first]!=top[v.first]){
        if(dep[top[u.first]]<dep[top[v.first]]) swap(u,v);
        Q temp=query(root,cnt[top[u.first]],cnt[u.first]+1);
        ans[u.second]=ans[u.second]+temp.sum-(lc[u.second]==temp.rcol);
        lc[u.second]=temp.lcol;
        u.first=fa[top[u.first]];
    }
    if(cnt[u.first]>cnt[v.first]) swap(u,v);
    Q te=query(root,cnt[u.first],cnt[v.first]+1);
    int anssum=ans[0]+ans[1]+te.sum-(lc[u.second]==te.lcol)-(lc[v.second]==te.rcol);
    return anssum;
}
int main(){
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++)
        scanf("%d",&C[i]);
    for(int i=1;i<n;i++){
        int a,b;
        scanf("%d%d",&a,&b);
        e[a].push_back(b);
        e[b].push_back(a);
    }
    dfs1(1,-1,0);
    dfs2(1,1);
    root=build(0,T);
    for(int i=1;i<=m;i++){
        char a[2];
        scanf("%s",a);
        if(a[0]=='C'){
            int a,b,c;
            scanf("%d%d%d",&a,&b,&c);
            changeans(a,b,c);
        }
        else if(a[0]=='Q'){
            int a,b;
            scanf("%d%d",&a,&b);
            printf("%d\n",getans(a,b));
        }
    }
    return 0;
}