洛谷 P2486 [SDOI2011]染色 LCT

Code:

#include <cstdio>          //SDOI2010 染色
#include <algorithm>
#include <cstring>
#include <string>
using namespace std;

void setIO(string a)
{
    freopen((a+".in").c_str(),"r",stdin);
}
struct LCT
{
    #define maxn 200000
    #define lson(x) ch[x][0]
    #define rson(x) ch[x][1]
    int ch[maxn][2],f[maxn],siz[maxn];
    int val[maxn],left[maxn],right[maxn];
    int tag[maxn];
    int sta[maxn];
    int lazy[maxn];
    int isRoot(int x)
    {
        return !(ch[f[x]][0]==x||ch[f[x]][1]==x);
    }
    int get(int x)
    {
        return ch[f[x]][1]==x;
    }
    void update(int x,int c)
    {
        if(!x)return;
        val[x]=left[x]=right[x]=c,siz[x]=1,lazy[x]=c;
    }
    void pushup(int x)
    {
        if(!x) return;
        siz[x]=siz[lson(x)]+siz[rson(x)]+1;
        left[x]=lson(x)?left[lson(x)]:val[x];
        right[x]=rson(x)?right[rson(x)]:val[x];
        if(right[lson(x)]==val[x]) siz[x]-=1;
        if(left[rson(x)]==val[x]) siz[x]-=1;
    }
    void mark(int x)
    {
        if(!x) return;
        swap(ch[x][0],ch[x][1]);
        swap(left[x],right[x]);
        tag[x]^=1;
    }
    void pushdown(int x)
    {
        if(lazy[x]) update(lson(x),lazy[x]),update(rson(x),lazy[x]),lazy[x]=0;
        if(tag[x]) mark(lson(x)),mark(rson(x)),tag[x]=0;
    }
    void rotate(int x)
    {
        int old=f[x], fold=f[old], which=get(x);
        if(!isRoot(old)) ch[fold][ch[fold][1]==old]=x;
        ch[old][which]=ch[x][which^1],f[ch[old][which]]=old;
        ch[x][which^1]=old,f[old]=x,f[x]=fold;
        pushup(old),pushup(x);
    }
    void splay(int x)
    {
        int v=0,u=x;
        sta[++v]=u;
        while(!isRoot(u)) sta[++v]=f[u],u=f[u];
        while(v) pushdown(sta[v--]);
        u=f[u];
        for(int fa;(fa=f[x])!=u;rotate(x))
            if(f[fa]!=u) rotate(get(fa)==get(x)?fa:x);
    }
    void Access(int x)
    {
        for(int y=0;x;y=x,x=f[x]) splay(x), ch[x][1]=y, pushup(x);
    }
    void makeRoot(int x)
    {
        Access(x), splay(x), mark(x);
    }
    void link(int a,int b)
    {
        makeRoot(a), f[a]=b;
    }
    void split(int a,int b)
    {
        makeRoot(a), Access(b), splay(b);
    }
}tree;
int main()
{
    //setIO("input");
    int n,m;
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;++i)
        scanf("%d",&tree.val[i]),tree.left[i]=tree.right[i]=tree.val[i],tree.siz[i]=1;
    for(int i=1;i<n;++i)
    {
        int a,b;
        scanf("%d%d",&a,&b);
        tree.link(a,b);
    }
    while(m--)
    {
        int x,y,z;
        char opt[20];
        scanf("%s",opt);
        switch(opt[0])
        {
            case  'C':
            {
                scanf("%d%d%d",&x,&y,&z);
                tree.split(x,y),tree.update(y,z);
                break;
            }
            case 'Q':
            {
                scanf("%d%d",&x,&y);
                tree.split(x,y), printf("%d\n",tree.siz[y]);
                break;
            }
        }
    }
    return 0;
}