HYSBZ 2243 染色

树上进行区间合并的操作，注意从当前链转到另外一条树链时，要考虑其节点与父亲节点的颜色是否相同

#include<stdio.h>
#include<vector>
#include<string.h>
#include<iostream>
using namespace std;
const int M=100010;
#define lson l,m,i*2
#define rson m+1,r,i*2+1
struct Edge
{
    int to,next;
}edge[M*2];
int head[M],tot;
int fa[M],son[M],top[M],num[M],deep[M],p[M],fp[M],pos;
int color[M];
void init()
{
    tot=pos=0;
    memset(head,-1,sizeof(head));
    memset(son,-1,sizeof(son));
}
void add(int u,int v)
{
    edge[tot].to=v;
    edge[tot].next=head[u];
    head[u]=tot++;
}
void dfs(int u,int pre,int d)
{
    num[u]=1;fa[u]=pre;deep[u]=d;
    for(int i=head[u];i!=-1;i=edge[i].next)
    {
        int v=edge[i].to;
        if(v!=pre)
        {
            dfs(v,u,d+1);
            num[u]+=num[v];
            if(son[u]==-1||num[v]>num[son[u]])
                son[u]=v;
        }
    }
}
void getpos(int u,int sp)
{
    top[u]=sp;
    p[u]=pos++;
    fp[p[u]]=u;
    if(son[u]==-1) return ;
    getpos(son[u],sp);
    for(int i=head[u];i!=-1;i=edge[i].next)
    {
        int v=edge[i].to;
        if(v!=fa[u]&&v!=son[u])
            getpos(v,v);
    }
}
struct Node
{
    int l,r,num,flag;
}node[M*3];
void push_up(int i)
{
    if(node[i*2].r==node[i*2+1].l)
        node[i].num=node[i*2].num+node[i*2+1].num-1;
    else
        node[i].num=node[i*2].num+node[i*2+1].num;
    node[i].l=node[i*2].l;
    node[i].r=node[i*2+1].r;
}
void push_down(int i)
{
    if(node[i].flag!=-1)
    {
        node[i*2].l=node[i*2+1].l=node[i].flag;
        node[i*2].r=node[i*2+1].r=node[i].flag;
        node[i*2].num=node[i*2+1].num=1;
        node[i*2].flag=node[i*2+1].flag=node[i].flag;
        node[i].flag=-1;
    }
}
void build(int l,int r,int i)
{
    node[i].flag=-1;
    node[i].l=node[i].r=color[fp[l]];
    node[i].num=1;
    if(l==r)
    {
        //cout<<" u="<<fp[l]<<" i="<<i<<" color="<<node[i].l<<endl;
        return ;
    }
    int m=(l+r)>>1;
    build(lson);
    build(rson);
    push_up(i);
}
void update(int l,int r,int i,int L,int R,int c)
{
    if(l>=L&&R>=r)
    {
        node[i].l=node[i].r=c;
        node[i].num=1;node[i].flag=c;
        return ;
    }
    if(L>r||R<l) return ;
    push_down(i);
    int m=(l+r)>>1;
    update(lson,L,R,c);
    update(rson,L,R,c);
    push_up(i);
}
int query(int l,int r,int i,int L,int R)
{
    if(L<=l&&r<=R)
    {
        return node[i].num;
    }
    push_down(i);
    int m=(l+r)>>1;
    if(R<=m) {return query(lson,L,R);}
    else if(L>m) {return query(rson,L,R);}
    else
    {
        int ans=query(lson,L,m)+query(rson,m+1,R);
        if(node[i*2].r==node[i*2+1].l) ans--;
        return ans;
    }
}
int query2(int l,int r,int i,int k)
{
    if(l>=r) return node[i].l;
    int m=(l+r)>>1;
    push_down(i);
    if(k<=m) return query2(lson,k);
    else return query2(rson,k);
}
void treeupdate(int u,int v,int c)
{
    int f1=top[u],f2=top[v];
    while(f1!=f2)
    {
        if(deep[f1]<deep[f2])
        {
            swap(f1,f2);swap(u,v);
        }
        update(0,pos-1,1,p[f1],p[u],c);
        u=fa[f1];f1=top[u];
    }
    if(deep[u]>deep[v]) swap(u,v);
    update(0,pos-1,1,p[u],p[v],c);
}
int findnum(int u,int v)
{
    int f1=top[u],f2=top[v],num=0;
    //cout<<f1<<" "<<f2<<" "<<u<<" "<<v<<endl;
    while(f1!=f2)
    {
        if(deep[f1]<deep[f2])
        {
            swap(f1,f2);swap(u,v);
        }
        int t=query(0,pos-1,1,p[f1],p[u]);
        num+=t;
        //cout<<"t="<<t<<" f1="<<f1<<" u="<<u<<" p[f1]="<<p[f1]<<" p[u]"<<p[u]<<endl;
        //if(query2(0,pos-1,1,p[f1])==query2(0,pos-1,1,p[fa[f1]])) num--;
        u=fa[f1];f1=top[u];
    }
    if(deep[u]>deep[v]) swap(u,v);
    int t=query(0,pos-1,1,p[u],p[v]);
    //cout<<"t="<<t<<" u="<<u<<" v="<<v<<endl;
    num+=t;
    return num;
}
int main()
{
    int n,m;
    while(scanf("%d%d",&n,&m)!=EOF)
    {
        for(int i=1;i<=n;i++)
            scanf("%d",&color[i]);
        init();
        for(int i=0;i<n-1;i++)
        {
            int u,v;
            scanf("%d%d",&u,&v);
            add(u,v);add(v,u);
        }
        dfs(1,0,0);
        getpos(1,1);
        build(0,pos-1,1);
        char op[10];
        int a,b,c;
        for(int i=0;i<m;i++)
        {
            scanf("%s",op);
            //printf("%s\n",op);
            if(op[0]=='C')
            {
                scanf("%d%d%d",&a,&b,&c);
                treeupdate(a,b,c);
            }
            else
            {
                scanf("%d%d",&a,&b);
                //cout<<a<<" "<<b<<endl;
                printf("%d\n",findnum(a,b));
            }
        }
    }
    return 0;
}