BZOJ2243 染色

树链剖分点更新，线段树区间更新，将某一段更新值为某个数C(颜色)，区间查询，查询区间内有多少颜色段

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
const int maxn = 100005;
int siz[maxn],son[maxn],top[maxn],tid[maxn],fa[maxn],dep[maxn];
int sum[maxn<<2],cl[maxn<<2],cr[maxn<<2],lz[maxn<<2];
struct node
{
    int v,next;
}e[maxn*2];
int head[maxn],val[maxn],cnt,lable;
int n,m;
void init()
{
    memset(head,-1,sizeof(head));
    cnt = lable = 0;
}
void add(int u,int v)
{
    e[cnt].v = v;
    e[cnt].next = head[u];
    head[u] = cnt++;
}
void find_heavy(int rt,int father,int depth)
{
    int maxsize = 0;
    fa[rt] = father;
    son[rt] = 0;
    siz[rt] = 1;
    dep[rt] = depth;
    for(int i = head[rt];i!=-1;i = e[i].next)if(e[i].v!=father){
        find_heavy(e[i].v,rt,depth+1);
        siz[rt]+=siz[e[i].v];
        if(siz[e[i].v]>maxsize)
            maxsize = siz[e[i].v],son[rt] = e[i].v;
    }
}
void connect(int rt,int anc)
{
    tid[rt] = ++lable;
    top[rt] = anc;
    if(son[rt])connect(son[rt],anc);
    for(int i = head[rt];i!=-1;i =e[i].next)
        if(e[i].v!=fa[rt]&&e[i].v!=son[rt])
            connect(e[i].v,e[i].v);
}
void pushup(int rt)
{
    sum[rt] = sum[rt<<1]+sum[rt<<1|1];
    cl[rt] = cl[rt<<1];
    cr[rt] = cr[rt<<1|1];
    if(cr[rt<<1]==cl[rt<<1|1])sum[rt]--;
}
void pushdown(int rt)
{
    if(lz[rt]!=-1){
        sum[rt<<1] = sum[rt<<1|1] = 1;
        cl[rt<<1] = cr[rt<<1] = lz[rt];
        cl[rt<<1|1] = cr[rt<<1|1] = lz[rt];
        lz[rt<<1] = lz[rt<<1|1] = lz[rt];
        lz[rt] = -1;
    }
}
void update(int L,int R,int c,int l,int r,int rt)
{
    if(L<=l&&r<=R){
        cl[rt] = c;
        cr[rt] = c;
        sum[rt] = 1;
        lz[rt] = c;
        return;
    }
    pushdown(rt);
    int m = (l+r)>>1;
    if(L<=m)update(L,R,c,lson);
    if(m<R)update(L,R,c,rson);
    pushup(rt);
}
int query(int L,int R,int l,int r,int rt)
{
    if(L<=l&&r<=R)return sum[rt];
    pushdown(rt);
    int m = (l+r)>>1,ret = 0,ok = -1;
    if(L<=m)ret+=query(L,R,lson),ok++;
    if(m<R)ret+=query(L,R,rson),ok++;
    if(ok==1){
        if(cr[rt<<1]==cl[rt<<1|1])ret--;
    }
    return ret;
}

int get(int pos,int l,int r,int rt)
{
    if(l==r)return cl[rt];
    pushdown(rt);
    int m = (l+r)>>1;
    if(pos<=m)return get(pos,lson);
    else return get(pos,rson);
}
int getsum(int x,int y)
{
    int ans = 0;
    while(top[x]!=top[y])
    {
        if(dep[top[x]]<dep[top[y]])swap(x,y);
        ans+=query(tid[top[x]],tid[x],1,n,1);
        if(get(tid[top[x]],1,n,1)==get(tid[fa[top[x]]],1,n,1))ans--;
        x = fa[top[x]];
    }
    if(dep[x]>dep[y])swap(x,y);
    ans+=query(tid[x],tid[y],1,n,1);
    return ans;
}
void change(int x,int y,int c)
{
    while(top[x]!=top[y])
    {
        if(dep[top[x]]<dep[top[y]])swap(x,y);
        update(tid[top[x]],tid[x],c,1,n,1);
        x = fa[top[x]];
    }
    if(dep[x]>dep[y])swap(x,y);
    update(tid[x],tid[y],c,1,n,1);
}
int main()
{
   // freopen("in.txt","r",stdin);
    while(~scanf("%d%d",&n,&m))
    {
        init();
        for(int i = 1;i<=n;++i)scanf("%d",val+i);
        for(int i = 1;i<n;++i){
            int u,v;scanf("%d%d",&u,&v);
            add(u,v);add(v,u);
        }
        find_heavy(1,1,1);
        connect(1,1);
        memset(lz,-1,sizeof(lz));
        memset(sum,0,sizeof(sum));
        for(int i = 1;i<=n;++i)update(tid[i],tid[i],val[i],1,n,1);
        while(m--)
        {
            char s[2];int a,b,c;
            scanf("%s",s);
            if(s[0]=='C'){
                scanf("%d%d%d",&a,&b,&c);
                change(a,b,c);
            }
            else {
                scanf("%d%d",&a,&b);
                printf("%d\n",getsum(a,b));
            }
        }
        puts("");
    }
    return 0;
}