【BZOJ 2243】染色 - 树链剖分+线段树

#include <cstdio>
#include <cstring>
#include <cstdlib>
using namespace std;

const int N=100001;

int n,m,color[N];   //Basic
struct G
{
    int v,nxt;
}map[N+N];
int hd[N],tt;   //Graph
int son[N],dep[N],pre[N],size[N];
int tid[N],top[N],num;  //Chain
struct T
{
    int l,r;
    int lc,rc,tag;
    int cnt;
}tr[N<<2];    //Segment Tree

inline int read(void)
{
    int s=0,f=1; char c=getchar();
    for (;c<'0'||c>'9';c=getchar()) if (c=='-') f=-1;
    for (;'0'<=c&&c<='9';c=getchar()) s=(s<<1)+(s<<3)+c-'0';
    return s*f;
}

inline void ins(int u,int v)
{
    map[++tt].v=v;
    map[tt].nxt=hd[u];
    hd[u]=tt;
}

void find(int now,int ht)
{
    dep[now]=ht,size[now]=1;
    for (int k=hd[now];k;k=map[k].nxt)
        if (!size[map[k].v])
        {
            find(map[k].v,ht+1);
            size[now]+=size[map[k].v];
            pre[map[k].v]=now;
            if (!son[now]||size[son[now]]<size[map[k].v]) son[now]=map[k].v;
        }
}

void cut(int now,int anc)
{
    top[now]=anc,tid[now]=++num;
    if (son[now]) cut(son[now],anc);
    for (int k=hd[now];k;k=map[k].nxt)
        if (!tid[map[k].v]) cut(map[k].v,map[k].v);
}

inline void clear(int now)
{
    if (tr[now].tag==-1) return;
    tr[now<<1].tag=tr[now<<1|1].tag=tr[now].tag;
    tr[now<<1].lc=tr[now<<1].rc=tr[now].tag;
    tr[now<<1|1].lc=tr[now<<1|1].rc=tr[now].tag;
    tr[now<<1].cnt=tr[now<<1|1].cnt=1;
    tr[now].tag=-1;
}

void build(int now,int l,int r)
{
    tr[now].l=l;
    tr[now].r=r;
    tr[now].tag=-1;
    if (l^r)
    {
        int mid=l+r>>1;
        build(now<<1,l,mid);
        build(now<<1|1,mid+1,r);
    }
}

void ins(int now,int l,int r,int c)
{
    if (l<=tr[now].l&&tr[now].r<=r)
    {
        tr[now].tag=tr[now].lc=tr[now].rc=c;
        tr[now].cnt=1;
        return;
    }
    clear(now);

    int mid=tr[now].l+tr[now].r>>1;
    if (l<=mid) ins(now<<1,l,r,c);
    if (mid<r) ins(now<<1|1,l,r,c);

    tr[now].lc=tr[now<<1].lc;
    tr[now].rc=tr[now<<1|1].rc;
    tr[now].cnt=tr[now<<1].cnt+tr[now<<1|1].cnt-(tr[now<<1].rc==tr[now<<1|1].lc);
}

void init(void)
{
    n=read(),m=read();
    for (int i=1;i<=n;i++) color[i]=read();

    int u,v;
    for (int i=1;i<n;i++)
    {
        u=read(),v=read();
        ins(u,v),ins(v,u);
    }

    find(1,1);
    cut(1,1);

    build(1,1,n);
    for (int i=1;i<=n;i++) ins(1,tid[i],tid[i],color[i]);
}

inline int LCA(int x,int y)
{
    for (;top[x]^top[y];)
        if (dep[top[x]]>dep[top[y]])
            x=pre[top[x]];
        else y=pre[top[y]];
    return dep[x]<dep[y]?x:y;
}

int query(int now,int l,int r)
{
    if (l<=tr[now].l&&tr[now].r<=r) return tr[now].cnt;
    clear(now);
    int mid=tr[now].l+tr[now].r>>1,sum=0;
    if (l<=mid) sum+=query(now<<1,l,r);
    if (mid<r) sum+=query(now<<1|1,l,r);
    if (l<=mid&&mid<r) sum-=tr[now<<1].rc==tr[now<<1|1].lc;
    return sum;
}

int findcolor(int now,int loc)
{
    if (tr[now].l==tr[now].r) return tr[now].tag;
    clear(now);
    int mid=tr[now].l+tr[now].r>>1;
    return findcolor(now<<1|(loc<=mid?0:1),loc);
}

void work(void)
{
    char c; int x,y,z,p,res,uc,dc;
    for (int i=1;i<=m;i++)
    {
        scanf("\n%c",&c);
        if (c=='C')
        {
            x=read(),y=read(),z=read(),p=LCA(x,y);
            for (;top[x]^top[p];x=pre[top[x]]) ins(1,tid[top[x]],tid[x],z);
            for (;top[y]^top[p];y=pre[top[y]]) ins(1,tid[top[y]],tid[y],z);
            ins(1,tid[p],x^p?tid[x]:tid[y],z);
        }
        else
        {
            x=read(),y=read(),p=LCA(x,y);
            uc=dc=-1,res=0;
            for (;top[x]^top[p];x=pre[top[x]])
            {
                if (uc^-1) 
                {
                    dc=findcolor(1,tid[x]);
                    res-=uc==dc;
                }
                res+=query(1,tid[top[x]],tid[x]);
                uc=findcolor(1,tid[top[x]]);
            }
            dc=findcolor(1,tid[x]);
            res-=uc==dc;
            uc=dc=-1;
            for (;top[y]^top[p];y=pre[top[y]])
            {
                if (uc^-1)
                {
                    dc=findcolor(1,tid[y]);
                    res-=uc==dc;
                }
                res+=query(1,tid[top[y]],tid[y]);
                uc=findcolor(1,tid[top[y]]);
            }
            dc=findcolor(1,tid[y]);
            res-=uc==dc;
            res+=query(1,tid[p],x^p?tid[x]:tid[y]);
            printf("%d\n",res);
        }
    }
}

int main(void)
{   
    init();
    work();

    return 0;
}