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=-,int b=-,int c=-):l(a),r(b),sum(c){}
 };
 color merge(color a,color b){
     color c;
     if (a.sum==-) return b;
     if (b.sum==-) return a;
     if (a.r==b.l){
         c.sum=a.sum+b.sum-;
         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<<];
     int lazy[V<<];
     void pushup(int rt){
         tag[rt]=merge(tag[rt<<],tag[rt<<|]);
     }
     void pushdown(int rt){
         if (lazy[rt]){
             lazy[rt<<]=lazy[rt<<|]=lazy[rt];
             tag[rt<<].l=tag[rt<<].r=lazy[rt];
             tag[rt<<|].l=tag[rt<<|].r=lazy[rt];
             tag[rt<<].sum=tag[rt<<|].sum=;
             lazy[rt]=;
             return;
         }
     }
     void build(int l,int r,int rt){
         if (l==r){
             tag[rt].l=tag[rt].r=a[rk[l]];
             tag[rt].sum=;
             return;
         }
         int m=(l+r)/;
         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=;
             lazy[rt]=val;
             return;
         }
         pushdown(rt);
         int m=(l+r)/;
         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)/;
         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]=; dep[u]=dep[fa[u]]+; son[u]=;
     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,,n,);
         x=fa[top[x]];
     }
     if (dep[x]>dep[y]) swap(x,y);
     T.update(w[x],w[y],val,,n,);
 }
 void pt(color a){
     printf("%d %d %d\n",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],,n,),ans1);
             x=fa[top[x]];
         }
         else
         {
             ans2=merge(T.query(w[top[y]],w[y],,n,),ans2);
             y=fa[top[y]];
         }

     }
     if (dep[x]<dep[y]){
         ans=T.query(w[x],w[y],,n,);
         ans=merge(ans,ans2);
         swap(ans.l,ans.r);
         ans=merge(ans,ans1);
     }
     else
     {
         ans=T.query(w[y],w[x],,n,);
         ans=merge(ans,ans1);
         swap(ans.l,ans.r);
         ans=merge(ans,ans2);
     }
     printf("%d\n",ans.sum );
 }

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