【树链剖分】bzoj2243 [SDOI2011]染色

树链剖分模板题。线段树维护每个段中的颜色数、左端点颜色、右端点颜色。

pushup: col[rt]=col[rt<<1]+col[rt<<1|1]-(Rcol[rt<<1]==Lcol[rt<<1|1]);

在合并各个链的答案时，要注意若两头颜色相同，ans--。

【教训：树链剖分题在进行建立线段树树时，要注意下面代码中的标注部分。否则会WA】

Code:

#include<cstdio>
#include<algorithm>
using namespace std;
#define lson rt<<1,l,m
#define rson rt<<1|1,m+1,r
#define maxn 100001
int col[maxn<<2],L[maxn<<2],R[maxn<<2],delta[maxn<<2],n,m,Val[maxn],Num[maxn];
int V[maxn<<1],First[maxn<<1],Next[maxn<<1],All_Lcol,All_Rcol,en;
bool vis[maxn];
int top[maxn],son[maxn],tot,dep[maxn],fa[maxn],siz[maxn],Map[maxn];
inline void AddEdge(const int &UU,const int &VV){V[++en]=VV;Next[en]=First[UU];First[UU]=en;}
inline void pushup(const int &rt)
{
    L[rt]=L[rt<<1];
    R[rt]=R[rt<<1|1];
    col[rt]=col[rt<<1]+col[rt<<1|1]-(R[rt<<1]==L[rt<<1|1]);
}
inline void pushdown(const int &rt)
{
    if(delta[rt]!=-1)
      {
          L[rt<<1]=R[rt<<1]=L[rt<<1|1]=R[rt<<1|1]=delta[rt];
          delta[rt<<1]=delta[rt<<1|1]=delta[rt];
          delta[rt]=-1;
          col[rt<<1]=col[rt<<1|1]=1;
      }
}
void buildtree(int rt,int l,int r)
{
    delta[rt]=-1;
    if(l==r)
      {
          col[rt]=1;
          L[rt]=R[rt]=Val[Map[l]];//l是线段树中的新位置，
                                //Map[l]表示新位置对应的在树中的编号是谁 
        return;
      }
    int m=l+r>>1;
    buildtree(lson);
    buildtree(rson);
    pushup(rt);
}
void update(int ql,int qr,int v,int rt,int l,int r)
{
    if(ql<=l&&r<=qr)
      {
        delta[rt]=L[rt]=R[rt]=v;//更新当前结点的标记值 
        col[rt]=1;
        return;
      }
    pushdown(rt);//将该节点的标记下传到孩子们 
    int m=l+r>>1;
    if(ql<=m)
      update(ql,qr,v,lson);
    if(m<qr)
      update(ql,qr,v,rson);
    pushup(rt);
}
int query(int ql,int qr,int rt,int l,int r)
{
    if(ql==l)
      All_Lcol=L[rt];
    if(r==qr)
      All_Rcol=R[rt];
    if(ql<=l&&r<=qr)
      return col[rt];
    pushdown(rt);//将该节点的标记下传到孩子们 
    int m=l+r>>1,res=0,now1=-1,now2=-1;
    if(ql<=m)
      {
          res+=query(ql,qr,lson);
          now1=R[rt<<1];
      }
    if(m<qr)
      {
          res+=query(ql,qr,rson);
          now2=L[rt<<1|1];
      }
    res-=(now1==now2 && now1!=-1);
    return res;
}
inline void Update(int u,int v,int val)
{
    int f1=top[u],f2=top[v];
    while(f1!=f2)
      {
          if(dep[f1]<dep[f2])
            {
                swap(u,v);
                swap(f1,f2);
            }
          update(Num[f1],Num[u],val,1,1,n);
          u=fa[f1];
          f1=top[u];
      }
    if(dep[u]<dep[v])
      swap(u,v);
    update(Num[v],Num[u],val,1,1,n);
}
inline int Query(int u,int v)
{
    int f1=top[u],f2=top[v],ans=0,t1=-1,t2=-1;
    while(f1!=f2)
      {
          if(dep[f1]<dep[f2])
            {
                swap(u,v);
                swap(f1,f2);
                swap(t1,t2);
            }
          ans+=query(Num[f1],Num[u],1,1,n);
          ans-=(t1==All_Rcol);
          t1=All_Lcol;
          u=fa[f1];
          f1=top[u];
      }
    if(dep[u]<dep[v])
      {
          swap(u,v);
          swap(t1,t2);
      }
    ans+=query(Num[v],Num[u],1,1,n);
    ans-=((t1==All_Rcol&&t1!=-1)+(t2==All_Lcol&&t2!=-1));
    return ans;
}
void dfs1(int cur,int father,int depth)
{
    fa[cur]=father;
    dep[cur]=depth;
    siz[cur]=1;
    for(int i=First[cur];i;i=Next[i])
      if(!vis[V[i]])
        {
          vis[V[i]]=true;
          dfs1(V[i],cur,depth+1);
          siz[cur]+=siz[V[i]];
          if(siz[V[i]]>siz[son[cur]])
            son[cur]=V[i];
          vis[V[i]]=false;
        }
}
void dfs2(int cur)
{
    if(son[cur]&&!vis[son[cur]])
      {
          vis[son[cur]]=true;
          top[son[cur]]=top[cur];
          Num[son[cur]]=++tot;
          Map[tot]=son[cur];
          dfs2(son[cur]);
          vis[son[cur]]=false;
      }
    for(int i=First[cur];i;i=Next[i])
      if(son[cur]!=V[i]&&!vis[V[i]])
        {
          vis[V[i]]=true;
          top[V[i]]=V[i];
          Num[V[i]]=++tot;
          Map[tot]=V[i];
          dfs2(V[i]);
          vis[V[i]]=false;
        }
}
int Res,num;
char C;
char CH[20];
inline int getint()
{
    Res=0;
    C='*'; 
    while(C<'0'||C>'9')
      C=getchar();
    while(C>='0'&&C<='9')
      {
        Res=Res*10+(C-'0');
        C=getchar();
      }
    return Res;
}
inline void putint(int x)
{
    num=0;  
    while(x>0)
      CH[++num]=x%10,x/=10;
    while(num)   
      putchar(CH[num--]+48);
    putchar('
');
}
int main()
{
    int x,y,a,b,c;
    char op;
    n=getint();
    m=getint();
    for(int i=1;i<=n;i++)
      Val[i]=getint();
    for(int i=1;i<n;i++){x=getint();y=getint();AddEdge(x,y);AddEdge(y,x);}
    top[1]=1;
    Num[1]=++tot;
    Map[tot]=1;
    vis[1]=true;
    dfs1(1,0,1);
    dfs2(1);
    buildtree(1,1,n);
    for(int i=1;i<=m;i++)
      {
          getchar();
          op=getchar();
          a=getint();
          b=getint();
          if(op=='C')
            {
                c=getint();
                Update(a,b,c);
            }
          else
            putint(Query(a,b));
      }
    return 0;
}