BZOJ 2243 [SDOI2011]染色：树剖【维护路径上颜色段】

题目链接：http://www.lydsy.com/JudgeOnline/problem.php?id=2243

题意：

　　给定一棵有n个节点的无根树和m个操作，操作有2类：

　　1、将节点a到节点b路径上所有点都染成颜色c；

　　2、询问节点a到节点b路径上的颜色段数量（连续相同颜色被认为是同一段），如"112221"由3段组成："11"、"222"和"1"。

　　请你写一个程序依次完成这m个操作。

题解：

　　树剖。

　　线段树上维护区间颜色段，每个节点：

　　　　dat[]：当前区间的颜色段数量

　　　　st[]：修改标记

　　　　lc[] and rc[]：当前区间的最左端和最右端颜色

　　　　lson[] and rson[]：左右儿子编号

　　线段树部分见NOI 2007 项链工厂。

　　树剖query中，每次a = par[tp[a]]往上跳之前，先判断getcol(dfsx[tp[a]]) == getcol(dfsx[par[tp[a]]])

　　如果相等，上下两端合并，sum--。

AC Code:

 #include <iostream>
 #include <stdio.h>
 #include <string.h>
 #include <vector>
 #define MAX_N 1000005

 using namespace std;

 int n,m;
 int tot=;
 int cnt=;
 int c[MAX_N];
 int dat[MAX_N];
 int st[MAX_N];
 int lc[MAX_N];
 int rc[MAX_N];
 int lson[MAX_N];
 int rson[MAX_N];
 int par[MAX_N];
 int dep[MAX_N];
 int siz[MAX_N];
 int tp[MAX_N];
 int son[MAX_N];
 int dfsx[MAX_N];
 int idx[MAX_N];
 vector<int> edge[MAX_N];

 void read()
 {
     scanf("%d%d",&n,&m);
     for(int i=;i<=n;i++)
     {
         scanf("%d",&c[i]);
     }
     int a,b;
     for(int i=;i<n;i++)
     {
         scanf("%d%d",&a,&b);
         edge[a].push_back(b);
         edge[b].push_back(a);
     }
 }

 void push_down(int now)
 {
     if(st[now]!=-)
     {
         dat[lson[now]]=dat[rson[now]]=;
         st[lson[now]]=st[rson[now]]=st[now];
         lc[lson[now]]=rc[lson[now]]=lc[rson[now]]=rc[rson[now]]=st[now];
         st[now]=-;
     }
 }

 void push_up(int now)
 {
     dat[now]=;
     lc[now]=lc[lson[now]];
     rc[now]=rc[rson[now]];
     if(lson[now]) dat[now]+=dat[lson[now]];
     if(rson[now]) dat[now]+=dat[rson[now]];
     if(lson[now] && rson[now] && rc[lson[now]]==lc[rson[now]]) dat[now]--;
 }

 int build(int l,int r)
 {
     int rt=++tot;
     dat[rt]=;
     st[rt]=-;
     lc[rt]=c[idx[l]];
     rc[rt]=c[idx[r]];
     lson[rt]=rson[rt]=;
     if(l<r)
     {
         int mid=(l+r)>>;
         lson[rt]=build(l,mid);
         rson[rt]=build(mid+,r);
         push_up(rt);
     }
     return rt;
 }

 void update(int a,int b,int k,int l,int r,int x)
 {
     if(a<=l && r<=b)
     {
         dat[k]=;
         st[k]=lc[k]=rc[k]=x;
         return;
     }
     if(r<a || b<l) return;
     push_down(k);
     int mid=(l+r)>>;
     update(a,b,lson[k],l,mid,x);
     update(a,b,rson[k],mid+,r,x);
     push_up(k);
 }

 int query(int a,int b,int k,int l,int r)
 {
     if(a<=l && r<=b) return dat[k];
     if(r<a || b<l) return ;
     push_down(k);
     int mid=(l+r)>>;
     int ans=;
     if(a<=mid) ans+=query(a,b,lson[k],l,mid);
     if(b>mid) ans+=query(a,b,rson[k],mid+,r);
     if(a<=mid && b>mid && rc[lson[k]]==lc[rson[k]]) ans--;
     return ans;
 }

 int getcol(int p,int k,int l,int r)
 {
     if(l==r) return lc[k];
     push_down(k);
     int mid=(l+r)>>;
     if(p<=mid) return getcol(p,lson[k],l,mid);
     else return getcol(p,rson[k],mid+,r);
 }

 void dfs1(int now,int p,int d)
 {
     par[now]=p;
     dep[now]=d;
     siz[now]=;
     for(int i=;i<edge[now].size();i++)
     {
         int temp=edge[now][i];
         if(temp!=p)
         {
             dfs1(temp,now,d+);
             siz[now]+=siz[temp];
         }
     }
 }

 void dfs2(int now,int anc)
 {
     tp[now]=anc;
     son[now]=-;
     cnt++;
     idx[cnt]=now;
     dfsx[now]=cnt;
     for(int i=;i<edge[now].size();i++)
     {
         int temp=edge[now][i];
         if((son[now]==- || siz[temp]>siz[son[now]]) && temp!=par[now])
         {
             son[now]=temp;
         }
     }
     if(son[now]!=-) dfs2(son[now],anc);
     for(int i=;i<edge[now].size();i++)
     {
         int temp=edge[now][i];
         if(temp!=par[now] && temp!=son[now]) dfs2(temp,temp);
     }
 }

 void update_chain(int a,int b,int x)
 {
     while(tp[a]!=tp[b])
     {
         if(dep[tp[a]]<dep[tp[b]]) swap(a,b);
         update(dfsx[tp[a]],dfsx[a],,,n,x);
         a=par[tp[a]];
     }
     if(dep[a]<dep[b]) swap(a,b);
     update(dfsx[b],dfsx[a],,,n,x);
 }

 int query_chain(int a,int b)
 {
     int sum=;
     while(tp[a]!=tp[b])
     {
         if(dep[tp[a]]<dep[tp[b]]) swap(a,b);
         sum+=query(dfsx[tp[a]],dfsx[a],,,n);
         if(getcol(dfsx[tp[a]],,,n)==getcol(dfsx[par[tp[a]]],,,n)) sum--;
         a=par[tp[a]];
     }
     if(dep[a]<dep[b]) swap(a,b);
     sum+=query(dfsx[b],dfsx[a],,,n);
     return sum;
 }

 void work()
 {
     dfs1(,-,);
     dfs2(,);
     build(,n);
     char s[];
     int a,b,c;
     while(m--)
     {
         scanf("%s",s);
         if(s[]=='Q')
         {
             scanf("%d%d",&a,&b);
             printf("%d\n",query_chain(a,b));
         }
         else
         {
             scanf("%d%d%d",&a,&b,&c);
             update_chain(a,b,c);
         }
     }
 }

 int main()
 {
     read();
     work();
 }