BZOJ2243 (树链剖分+线段树）

Problem 染色（BZOJ2243）

题目大意

　　给定一颗树，每个节点上有一种颜色。

　　要求支持两种操作：

　　　　操作1：将a->b上所有点染成一种颜色。

　　　　操作2：询问a->b上的颜色段数量。

解题分析

　　树链剖分+线段树。

　　开一个记录类型，记录某一段区间的信息。l 表示区间最左侧的颜色 ， r 表示区间最右侧的颜色 ， sum 表示区间中颜色段数量。

　　合并时判断一下左区间的右端点和有区间的左端点的颜色是否一样。

　　树上合并时需要用两个变量ans1，ans2来存储。ans1表示x往上走时形成的链的信息，ans2表示y往上走时形成链的信息。

　　当x和y位于同一条重链上时，有三个区间需要合并在一起，注意判断顺序。

参考程序

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

#define V 100008
#define E 200008
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1

int n,m,cnt;
int a[V],size[V],dep[V],fa[V],son[V],top[V],w[V],rk[V];

struct line{
    int u,v,nt;
}eg[E];
int sum,lt[V];

struct color{
    int l,r,sum;
    color(int a=-1,int b=-1,int c=-1):l(a),r(b),sum(c){}
};
color merge(color a,color b){
    color c;
    if (a.sum==-1) return b;
    if (b.sum==-1) return a;
    if (a.r==b.l){
        c.sum=a.sum+b.sum-1;
        c.l=a.l;
        c.r=b.r;
    }
    else
    {
        c.sum=a.sum+b.sum;
        c.l=a.l;
        c.r=b.r;
    }
    return c;
}

struct segment_tree{
    color tag[V<<2];
    int lazy[V<<2];
    void pushup(int rt){
        tag[rt]=merge(tag[rt<<1],tag[rt<<1|1]);
    }
    void pushdown(int rt){
        if (lazy[rt]){
            lazy[rt<<1]=lazy[rt<<1|1]=lazy[rt];
            tag[rt<<1].l=tag[rt<<1].r=lazy[rt];
            tag[rt<<1|1].l=tag[rt<<1|1].r=lazy[rt];
            tag[rt<<1].sum=tag[rt<<1|1].sum=1;
            lazy[rt]=0;
            return;
        }
    }
    void build(int l,int r,int rt){
        if (l==r){
            tag[rt].l=tag[rt].r=a[rk[l]];
            tag[rt].sum=1;
            return;
        }
        int m=(l+r)/2;
        build(lson);
        build(rson);
        pushup(rt);
    }
    void update(int L,int R,int val,int l,int r,int rt){
        if (L<=l && r<=R){
            tag[rt].l=tag[rt].r=val;
            tag[rt].sum=1;
            lazy[rt]=val;
            return;
        }
        pushdown(rt);
        int m=(l+r)/2;
        if (L <= m) update(L,R,val,lson);
        if (m <  R) update(L,R,val,rson);
        pushup(rt);
    }
    color query(int L,int R,int l,int r,int rt){
        if (L<=l && r<=R){
            return tag[rt];
        }
        pushdown(rt);
        color res;
        int m=(l+r)/2;
        if (L <= m) res=merge(res,query(L,R,lson));
        if (m <  R) res=merge(res,query(L,R,rson));
        return res;
    }
}T;

void adt(int u,int v){
    eg[++sum].u=u; eg[sum].v=v; eg[sum].nt=lt[u]; lt[u]=sum;
}
void add(int u,int v){
    adt(u,v); adt(v,u);    
}

void dfs_1(int u){
    size[u]=1; dep[u]=dep[fa[u]]+1; son[u]=0;
    for (int i=lt[u];i;i=eg[i].nt){
        int v=eg[i].v;
        if (v==fa[u]) continue;
        fa[v]=u;
        dfs_1(v);
        size[u]+=size[v];
        if (size[v]>size[son[u]]) son[u]=v;
    }
}
void dfs_2(int u,int tp){
    w[u]=++cnt; top[u]=tp; rk[cnt]=u;
    if (son[u]) dfs_2(son[u],tp);
    for (int i=lt[u];i;i=eg[i].nt){
        int v=eg[i].v;
        if (v==fa[u] || v==son[u]) continue;
        dfs_2(v,v);
    }
}
void change(int x,int y,int val){
    while (top[x]!=top[y]){
        if (dep[top[x]]<dep[top[y]]) swap(x,y);
        T.update(w[top[x]],w[x],val,1,n,1);
        x=fa[top[x]];
    }
    if (dep[x]>dep[y]) swap(x,y);
    T.update(w[x],w[y],val,1,n,1);
}
void pt(color a){
    printf("%d %d %d
",a.l,a.r,a.sum);
}
void find(int x,int y){
    color ans1,ans2,ans;
    while (top[x]!=top[y]){
        if (dep[top[x]]>dep[top[y]]){
            ans1=merge(T.query(w[top[x]],w[x],1,n,1),ans1);
            x=fa[top[x]];
        }
        else
        {
            ans2=merge(T.query(w[top[y]],w[y],1,n,1),ans2);
            y=fa[top[y]];
        }
        
    }
    if (dep[x]<dep[y]){ 
        ans=T.query(w[x],w[y],1,n,1);
        ans=merge(ans,ans2);
        swap(ans.l,ans.r);
        ans=merge(ans,ans1);
    }
    else
    {
        ans=T.query(w[y],w[x],1,n,1);
        ans=merge(ans,ans1);
        swap(ans.l,ans.r);
        ans=merge(ans,ans2);
    }
    printf("%d
",ans.sum );
}

int main(){
    memset(lt,0,sizeof(lt)); sum=1; cnt=0;
    scanf("%d %d",&n,&m);
    for (int i=1;i<=n;i++) scanf("%d",&a[i]);
    for (int i=1;i<n;i++){
        int u,v;
        scanf("%d %d",&u,&v);
        add(u,v);
    }
    dfs_1(1);
    dfs_2(1,1);
    T.build(1,n,1);
    while (m--){
        char ch[2];
        int x,y,z;
        scanf("%s",ch);
        if (ch[0]=='Q'){
            scanf("%d %d",&x,&y);
            find(x,y);
        }
        else
        {    
            scanf("%d %d %d",&x,&y,&z);
            change(x,y,z);
        };
    }
}