bzoj2243 [SDOI2011]染色

　　树链剖分，线段树维护三个值:区间颜色数、区间左端颜色、区间右端颜色，当左区间右端颜色=右区间左端颜色时候，总颜色-1，树上的维护类似。

　　代码

#include<cstdio>
const int M = 1001010;
const int N = 300010;
int dp,p[N],pre[N],tt[N],size[N],f[N],gf[N],go[N],deep[N],color[N];
int id[N],n,m,i,a,b,c,tot,e[N];
int l[M],r[M],lc[M],rc[M],s[M],v[M]; 
void link(int x,int y)
{
    dp++;pre[dp]=p[x];p[x]=dp;tt[dp]=y;
}
struct O{
    int ans,lc,rc;
    void clear()
    {
        ans=0;lc=0;rc=0;
    }
    void set(int a,int b,int c)
    {
        ans=a;lc=b;rc=c;
    }
};
void dfs(int x,int fa)
{
    int i=p[x];
    size[x]=1;
    deep[x]=deep[fa]+1;
    while (i)
    {
        if (tt[i]!=fa)
        {
            f[tt[i]]=x;
            dfs(tt[i],x);
            if (size[tt[i]]>size[go[x]]) go[x]=tt[i];
            size[x]+=size[tt[i]]; 
        }
        i=pre[i];
    }
}
void gao(int x,int Fa,int fa)
{
    int i=p[x];
    gf[x]=Fa;id[x]=++tot;e[tot]=color[x];
    if (go[x]) gao(go[x],Fa,x);
    while (i)
    {
        if ((tt[i]!=fa)&&(tt[i]!=go[x]))
            gao(tt[i],tt[i],x);
        i=pre[i];
    }
}
void updata(int x)
{
    lc[x]=lc[2*x];
    rc[x]=rc[2*x+1];
    s[x]=s[2*x]+s[2*x+1];
    if (rc[2*x]==lc[2*x+1]) s[x]--;
}
void build(int x,int a,int b)
{
    int m;
    l[x]=a;r[x]=b;v[x]=-1;
    if (b-a>1)
    {
        m=(a+b)>>1;
        build(2*x,a,m);
        build(2*x+1,m,b);
        updata(x);
    }
    else
    {
        lc[x]=e[b];
        rc[x]=e[b];
        s[x]=1;
    }
}
void clean(int x)
{
    if (v[x]>=0)
    {
        s[x]=1;
        lc[x]=rc[x]=v[x];
        v[2*x]=v[x];
        v[2*x+1]=v[x];
        v[x]=-1;
    }
}
void change(int x,int a,int b,int c)
{
    int m;
    clean(x);
    if ((a<=l[x])&&(r[x]<=b))
    {
        v[x]=c;
        return;
    }
    m=(l[x]+r[x])>>1;
    if (a<m) change(2*x,a,b,c);
    if (m<b) change(2*x+1,a,b,c);
    clean(2*x);clean(2*x+1);
    updata(x);
}
O query(int x,int a,int b)
{
    int m;
    O tmp;tmp.clear();
    clean(x);
    if ((a<=l[x])&&(r[x]<=b))
    {
        tmp.set(s[x],lc[x],rc[x]);
        return tmp;
    }
    m=(l[x]+r[x])>>1;
    int L=-1,R=-1,A=-1;
    if (a<m) 
    {
        tmp=query(2*x,a,b);
        L=tmp.lc;
        R=tmp.rc;
        A=tmp.ans;    
    }
    if (m<b)
    {
        tmp=query(2*x+1,a,b);
        if (L==-1)
        {
            L=tmp.lc;
            R=tmp.rc;
            A=tmp.ans;    
        }
        else
        {
            A+=tmp.ans;
            if (R==tmp.lc) A--;
            R=tmp.rc;
        }
    }
    tmp.set(A,L,R);return tmp;
}
void gai(int a,int b,int c)
{
    while (1)
    {
        if (gf[a]==gf[b])
        {
            if (deep[a]<deep[b]) a^=b^=a^=b;
            change(1,id[b]-1,id[a],c);
            return;
        }
        else
        {
            if (deep[gf[a]]<deep[gf[b]]) a^=b^=a^=b;
            change(1,id[gf[a]]-1,id[a],c);
            a=f[gf[a]];
        }
    }
}
int wen(int a,int b)
{
    int ac=-1,bc=-1,ans=0;
    while (1)
    {
        if (gf[a]==gf[b])
        {
            if (deep[a]<deep[b]) a^=b^=a^=b,ac^=bc^=ac^=bc;
            O tmp=query(1,id[b]-1,id[a]);
            ans+=tmp.ans;
            if (tmp.rc==ac) ans--;
            if (tmp.lc==bc) ans--;
            return ans;
        }
        else
        {
            if (deep[gf[a]]<deep[gf[b]]) a^=b^=a^=b,ac^=bc^=ac^=bc;
            O tmp=query(1,id[gf[a]]-1,id[a]);
            ans+=tmp.ans;
            if (tmp.rc==ac) ans--;ac=tmp.lc;
            a=f[gf[a]];
        }
    }
}
int main()
{
    scanf("%d%d",&n,&m);
    for (i=1;i<=n;i++)
        scanf("%d",&color[i]);
    for (i=1;i<n;i++)
    {
        scanf("%d%d",&a,&b);
        link(a,b);
        link(b,a);
    }
    dfs(1,0);
    gao(1,1,0);
    build(1,0,n);
    for (i=1;i<=m;i++)
    {
        char typ;
        scanf(" %c",&typ);
        if (typ=='Q')
        {
            scanf("%d%d",&a,&b);
            printf("%d
",wen(a,b));
        }
        else
        {
            scanf("%d%d%d",&a,&b,&c);
            gai(a,b,c);
        }
    }
}