[BZOJ 2243][SDOI2011] 染色

可以说是树链剖分的模板题吧

基本思想就是先树链剖分，然后用线段树维护区间，区间的颜色种类个数，区间左端点的颜色，区间右端点的颜色

区间合并时，区间的颜色种类个数=左区间个数+右区间的个数，如果左区间的右端点的颜色等于右区间的左端点的颜色，答案还要减一

下面是代码：(用的是栈模拟递归)

#include<cmath>
#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<iostream>
#define N 100100
#define M 200100
using namespace std;

int n,m,a,b,p;
int color[N],head[N],go[M],next[M];

void build(int a,int b)
{
    go[++p]=b;
    next[p]=head[a];
    head[a]=p;
}

struct dfsdata
{
    int a,i;
    void newdata(int x,int y){a=x;i=y;}
}Stack[N];

struct makedata
{
    int a,i,yes;
    void newdata(int x,int y,int z=1){a=x;i=y;yes=z;}
}S[N];

int tail,f[N],size[N],max_son[N],u,v,deep[N];

void dfs()
{
    int i;
    Stack[tail=1].newdata(1,head[1]);f[1]=0;size[1]=1;deep[1]=1;
    while (tail)
    {
        i=Stack[tail].i;v=Stack[tail].a;
        if (!i)
        {
            size[f[v]]+=size[v];
            if (size[v]>size[max_son[f[v]]]) max_son[f[v]]=v;
            tail--;
            continue;
        }
        u=go[i];
        Stack[tail].i=next[i];
        if (f[v]!=u)
        {
            size[u]=1;
            f[u]=v;
            deep[u]=deep[v]+1;
            Stack[++tail].newdata(u,head[u]);
        }
    }
}

int L[N],l_deep[N],top[N],len,que[N];

void make()
{
    int i;
    S[tail=1].newdata(1,head[1]);L[1]=len=1;l_deep[1]=1;top[1]=1;que[1]=1;
    while (tail)
    {
        v=S[tail].a;i=S[tail].i;
        if (!i || size[v]==1){tail--;continue;}
        if (S[tail].yes)
        {
            S[tail].yes=0;
            u=max_son[v];
            L[u]=++len;
            que[len]=u;
            l_deep[u]=l_deep[v];
            top[u]=top[v];
            S[++tail].newdata(u,head[u]);
            continue;
        }
        u=go[i];
        S[tail].i=next[i];
        if (u!=f[v] && u!=max_son[v])
        {
            L[u]=++len;
            que[len]=u;
            l_deep[u]=l_deep[v]+1;
            top[u]=u;
            S[++tail].newdata(u,head[u]);
        }
    }
}

//==========================SegmentTree===========================//
int tree[N<<2],SL[N<<2],SR[N<<2],LZ[N<<2],lazy[N<<2];

void update(int a)
{
    tree[a]=tree[a<<1]+tree[(a<<1)+1];
    SL[a]=SL[a<<1];SR[a]=SR[(a<<1)+1];
    if (SR[a<<1]==SL[(a<<1)+1]) tree[a]--;
}

void BuildTree(int p,int l,int r)
{
    if (l==r)
    {
        tree[p]=1;
        SL[p]=SR[p]=color[que[l]];
        return;
    }
    int mid=(l+r)>>1;
    BuildTree(p<<1,l,mid);
    BuildTree((p<<1)+1,mid+1,r);
    update(p);
}

void Push_down(int a)
{
    if (!LZ[a]) return;
    LZ[a]=0;
    LZ[a<<1]=LZ[(a<<1)+1]=1;
    lazy[a<<1]=lazy[(a<<1)+1]=lazy[a];
    SL[a<<1]=SR[a<<1]=SL[(a<<1)+1]=SR[(a<<1)+1]=lazy[a];
    tree[a]=tree[a<<1]=tree[(a<<1)+1]=1;
}

void Change(int p,int l,int r,int fl,int fr,int c)
{
    if (fl<=l&&r<=fr)
    {
        LZ[p]=1;
        lazy[p]=c;
        SL[p]=SR[p]=c;
        tree[p]=1;
        return;
    }
    Push_down(p);
    int mid=(l+r)>>1;
    if (fl<=mid) Change(p<<1,l,mid,fl,fr,c);
    if (fr>mid) Change((p<<1)+1,mid+1,r,fl,fr,c);
    update(p);
}

int Query(int p,int l,int r,int fl,int fr,int al,int &x,int ar,int &y)
{
    if (l==fl&&r==fr)
    {
        if (LZ[p])
        {
            if (l==r)
            {
                LZ[p]=0;
                tree[p]=1;
                SL[p]=SR[p]=lazy[p];
            }
            else Push_down(p),update(p);
        }
        if (l==al) x=SL[p];
        if (r==ar) y=SR[p];
        return tree[p];
    }
    Push_down(p);
    int mid=(l+r)>>1;
    if (fr<=mid) return Query(p<<1,l,mid,fl,fr,al,x,ar,y);
    if (fl>mid) return Query((p<<1)+1,mid+1,r,fl,fr,al,x,ar,y);
    int ans=Query(p<<1,l,mid,fl,mid,al,x,ar,y)+Query((p<<1)+1,mid+1,r,mid+1,fr,al,x,ar,y);
    if (SR[p<<1]==SL[(p<<1)+1]) ans--;
    return ans;
}
//==========================SegmentTree===========================//

void Change(int a,int b,int c)
{
    while (top[a]!=top[b])
    {
        if (l_deep[a]<l_deep[b]) swap(a,b);
        color[a]=c;
        if (top[a]!=a)
        {
            Change(1,1,n,L[top[a]],L[a],c);
            a=top[a];
        }
        else
        {
            Change(1,1,n,L[a],L[a],c);
            a=f[a];
        }
    }
    color[a]=color[b]=c;
    if (L[a]>L[b]) swap(a,b);
    Change(1,1,n,L[a],L[b],c);
}

void Query(int a,int b)
{
    int ans=0,ac=-1,bc=-1,x,y;
    if (a==b)
    {
        printf("1\n");
        return;
    }
    while (top[a]!=top[b])
    {
        if (l_deep[a]<l_deep[b]) swap(a,b),swap(ac,bc);
        if (top[a]!=a)
        {
            ans+=Query(1,1,n,L[top[a]]+1,L[a],L[top[a]]+1,x,L[a],y);
            if (y==ac) ans--;
            ac=x;
            a=top[a];
        }
        else
        {
            ans+=Query(1,1,n,L[a],L[a],L[a],x,L[a],y);
            if (y==ac) ans--;
            ac=x;
            a=f[a];
        }
    }
    if (L[a]>L[b]) swap(a,b),swap(ac,bc);
    ans+=Query(1,1,n,L[a],L[b],L[a],x,L[b],y);
    if (x==ac) ans--;
    if (y==bc) ans--;
    printf("%d\n",ans);
}

int main()
{
    scanf("%d%d",&n,&m);
    for (int i=1;i<=n;i++) scanf("%d",color+i);
    for (int i=1;i<n;i++)
    {
        scanf("%d%d\n",&a,&b);
        build(a,b);
        build(b,a);
    }
    dfs();
    make();
    BuildTree(1,1,n);
    for (int i=1;i<=m;i++)
    {
        char ch;int c;
        scanf("%c",&ch);
        if (ch=='C')
        {
            scanf("%d%d%d\n",&a,&b,&c);
            Change(a,b,c);
        }
        else
        {
            scanf("%d%d\n",&a,&b);
            Query(a,b);
        }
    }
    return 0;
}