bzoj2243 SDOI2011 染色 树链剖分

树剖裸题，复习一波。。一开始还看错题目了以为要维护一段区间的颜色数量。。还以为要上主席树。。

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<iostream>
#include<queue>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
using namespace std;
const int N=5e5+5;
int n,m,cnt;
int val[N],head[N],next[N],go[N],fa[N],pos[N];
int a[N],tot;
int siz[N],son[N],top[N],dep[N];
struct node
{
    int l,r,sum,lc,rc,lazy;
}t[N*10];
inline void add(int x,int y)
{
    go[++tot]=y;
    next[tot]=head[x];
    head[x]=tot;
}
inline void dfs1(int x,int f)
{
    siz[x]=1;
    for(int i=head[x];i;i=next[i])
    {
        int v=go[i];
        if (v!=f)
        {
            dep[v]=dep[x]+1;
            fa[v]=x;
            dfs1(v,x);
            siz[x]+=siz[v];
            if (!son[x]||siz[son[x]]<siz[v])
            son[x]=v;
        }
    }
}
int tim;
inline void dfs2(int x,int chain)
{
    int k=0;
    pos[x]=++tim;
    top[x]=chain;
    for(int i=head[x];i;i=next[i])
    {
        int v=go[i];
        if (v!=fa[x])
        {
            if (siz[v]>siz[k])
            k=v;
        }
    }
    if (!k)return ;
    dfs2(k,chain);
    for(int i=head[x];i;i=next[i])
    {
        int v=go[i];
        if (v!=fa[x]&&v!=k)
        dfs2(v,v);
    }
}
inline void build(int x,int l,int r)
{
    t[x].l=l,t[x].r=r;
    if (l==r)return ;
    int mid=(l+r)>>1;
    build(x*2,l,mid);
    build(x*2+1,mid+1,r);
}
inline void pushlazy(int x)
{
    if (t[x].lazy==0||t[x].l==t[x].r)return;
    int v=t[x].lazy;
    t[x].lazy=0;
    t[x<<1].sum=t[x<<1|1].sum=1;
    t[x<<1].lc=t[x<<1].rc=t[x<<1|1].lc=t[x<<1|1].rc=t[x<<1].lazy=t[x<<1|1].lazy=v;
}
inline void queryc(int x,int l,int r,int v)
{
    if (t[x].l==l&&t[x].r==r)
    {
        t[x].sum=1;
        t[x].lc=t[x].rc=t[x].lazy=v;
        return;
    }
    pushlazy(x);
    int mid=(t[x].l+t[x].r)>>1;
    if (r<=mid)queryc(x<<1,l,r,v);
    else if (l>mid)queryc(x<<1|1,l,r,v);
    else
    {
        queryc(x<<1,l,mid,v);
        queryc(x<<1|1,mid+1,r,v);
    }
    t[x].lc=t[x<<1].lc;
    t[x].rc=t[x<<1|1].rc;
    t[x].sum=t[x<<1].sum+t[x<<1|1].sum;
    if (t[x<<1].rc==t[x<<1|1].lc)t[x].sum--;
}
inline void solvec(int x,int y,int z)
{
    while (top[x]!=top[y])
    {
        if (dep[top[x]]<dep[top[y]])swap(x,y);
        queryc(1,pos[top[x]],pos[x],z);
        x=fa[top[x]];
    }
    if (dep[x]>dep[y])swap(x,y);
    queryc(1,pos[x],pos[y],z);
}
inline int queryq(int x,int l,int r)
{
    pushlazy(x);
    if (t[x].l==l&&t[x].r==r)return t[x].sum;
    int mid=(t[x].l+t[x].r)>>1;
    if (r<=mid)return queryq(x<<1,l,r);
    else if (l>mid)return queryq(x<<1|1,l,r);
    else
    {
        int sum=queryq(x<<1,l,mid)+queryq(x<<1|1,mid+1,r);
        if (t[x<<1].rc==t[x<<1|1].lc)sum--;
        return sum;
    }
}
inline int find(int x,int pos)
{
    pushlazy(x);
    if (t[x].sum==1)return t[x].lc;
    int mid=(t[x].l+t[x].r)>>1;
    if (pos<=mid)return find(x<<1,pos);
    else find(x<<1|1,pos);
}
inline int solveq(int x,int y)
{
    int sum=0;
    while (top[x]!=top[y])
    {
        if (dep[top[x]]<dep[top[y]])swap(x,y);
        sum+=queryq(1,pos[top[x]],pos[x]);
        if (find(1,pos[top[x]])==find(1,pos[fa[top[x]]]))sum--;//注意合并的时候相邻颜色相同sum-1 
        x=fa[top[x]];
    }
        if (dep[x]>dep[y])swap(x,y);
    sum+=queryq(1,pos[x],pos[y]);
    return sum;
}
int main()
{
    scanf("%d%d",&n,&m);
    fo(i,1,n)scanf("%d",&a[i]);
    fo(i,1,n-1)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y);
        add(y,x);
    }
    dep[1]=1;
    dfs1(1,0);
    dfs2(1,1);
    build(1,1,n);
    char ch[2];
    fo(i,1,n)
    queryc(1,pos[i],pos[i],a[i]);
    fo(i,1,m)
    {
        scanf("%s",ch);
        if (ch[0]=='C')
        {
            int x,y,z;
            scanf("%d%d%d",&x,&y,&z);
            solvec(x,y,z);
        }
        else
        {
            int x,y;
            scanf("%d%d",&x,&y);
            printf("%d\n",solveq(x,y));
        }
    }
}