[BZOJ]2243: [SDOI2011]染色

#include<iostream>
#include<cstdio>
#include<cstring>
#define maxn 100009
using namespace std;

int n,m,cnt,sumedge;
int head[maxn],dad[maxn],size[maxn],deep[maxn],top[maxn];
int c[maxn],re[maxn],tpos[maxn];
char s[10];

struct Tree{
    int l,r,sum,lc,rc,s;
}tr[maxn<<2];

struct Edge{
    int x,y,nxt;
    Edge(int x=0,int y=0,int nxt=0):
        x(x),y(y),nxt(nxt){}
}edge[maxn<<1];

void add(int x,int y){
    edge[++sumedge]=Edge(x,y,head[x]);
    head[x]=sumedge;
}

void dfs(int x){
    size[x]=1;deep[x]=deep[dad[x]]+1;
    for(int i=head[x];i;i=edge[i].nxt){
        int v=edge[i].y;
        if(v==dad[x])continue;
        dad[v]=x;dfs(v);
        size[x]+=size[v];
    }
}

void dfs_(int x){
    int s=0;tpos[x]=++cnt;re[cnt]=x;
    if(!top[x])top[x]=x;
    for(int i=head[x];i;i=edge[i].nxt){
        int v=edge[i].y;
        if(v!=dad[x]&&size[v]>size[s])s=v;
    }
    if(s){
        top[s]=top[x];
        dfs_(s);
    }
    for(int i=head[x];i;i=edge[i].nxt){
        int v=edge[i].y;
        if(v!=dad[x]&&v!=s)dfs_(v);
    }
}

void pushup(int rt){
    tr[rt].sum=tr[rt<<1].sum+tr[rt<<1|1].sum-(tr[rt<<1].rc==tr[rt<<1|1].lc);
    tr[rt].lc=tr[rt<<1].lc;tr[rt].rc=tr[rt<<1|1].rc;
    return;
}

void pushdown(int rt){
    if(!tr[rt].s)return;
    tr[rt<<1].sum=tr[rt<<1|1].sum=1;
    tr[rt<<1].rc=tr[rt<<1].lc=tr[rt<<1|1].lc=tr[rt<<1|1].rc=tr[rt].s;
    tr[rt<<1].s=tr[rt<<1|1].s=tr[rt].s;
    tr[rt].s=0;return;
}

void build(int rt,int l,int r){
    tr[rt].l=l;tr[rt].r=r;
    if(l==r){
        tr[rt].sum=1;tr[rt].lc=tr[rt].rc=c[re[l]];
        return;
    }
    int mid=(l+r)>>1;
    build(rt<<1,l,mid);build(rt<<1|1,mid+1,r);
    pushup(rt);
}

void change(int rt,int l,int r,int ql,int qr,int p){
    if(l>=ql&&r<=qr){
        tr[rt].sum=1;tr[rt].lc=tr[rt].rc=p;tr[rt].s=p;
        return;
    }
    pushdown(rt);
    int mid=(l+r)>>1;
    if(ql<=mid)change(rt<<1,l,mid,ql,qr,p);
    if(qr>mid)change(rt<<1|1,mid+1,r,ql,qr,p);
    pushup(rt);
}

int query_sum(int rt,int l,int r,int ql,int qr){
    if(l>=ql&&r<=qr){
        return tr[rt].sum;
    }
    pushdown(rt);
    int mid=(l+r)>>1;
    if(qr<=mid)return query_sum(rt<<1,l,mid,ql,qr);
    if(ql>mid)return query_sum(rt<<1|1,mid+1,r,ql,qr);
    return query_sum(rt<<1,l,mid,ql,qr)+query_sum(rt<<1|1,mid+1,r,ql,qr)-(tr[rt<<1].rc==tr[rt<<1|1].lc);//**
}

int getcolr(int rt,int l,int r,int p){
    if(l==r)return tr[rt].lc;
    pushdown(rt);
    int mid=(l+r)>>1;
    if(p<=mid)return getcolr(rt<<1,l,mid,p);
    if(p>mid)return getcolr(rt<<1|1,mid+1,r,p);
}

int query(int x,int y){
    int ret=0;
    for(;top[x]!=top[y];){
        if(deep[top[x]]>deep[top[y]])swap(x,y);
        ret+=query_sum(1,1,n,tpos[top[y]],tpos[y]);
        if(getcolr(1,1,n,tpos[top[y]])==getcolr(1,1,n,tpos[dad[top[y]]]))ret--;
        y=dad[top[y]];
    }
    if(deep[x]>deep[y])swap(x,y);
    ret+=query_sum(1,1,n,tpos[x],tpos[y]);
    return ret;
}

void change_(int x,int y,int c){
    for(;top[x]!=top[y];){
        if(deep[top[x]]>deep[top[y]])swap(x,y);
        change(1,1,n,tpos[top[y]],tpos[y],c);
        y=dad[top[y]];
    }
    if(deep[x]>deep[y])swap(x,y);
    change(1,1,n,tpos[x],tpos[y],c);
}

int main(){
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++)scanf("%d",&c[i]);
    for(int i=1;i<n;i++){
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y);add(y,x);
    }
    dfs(1);dfs_(1);build(1,1,n);
    for(int i=1;i<=m;i++){
        int x,y,z;
        scanf("%s",s);
        if(s[0]==‘Q‘){
            scanf("%d%d",&x,&y);
            printf("%d\n",query(x,y));
        }else  {
            scanf("%d%d%d",&x,&y,&z);
            change_(x,y,z);
        }
    }
    return 0;
}