bzoj 2243: [SDOI2011]染色 （树链剖分+线段树 区间合并）

#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
#define mid ll m = (l+r)>>1
const ll M = 1e5+10;
ll cnt,n,q;
ll siz[M],son[M],fa[M],top[M],rk[M],tid[M],dep[M],a[M],cnt1,head[M],lazy[M<<2];
struct node{ll to,next;}e[M];
struct node1{ll ls,rs,val;};
node1 sum[M<<2];

void add(ll u,ll v){
    e[++cnt1].to = v;e[cnt1].next = head[u];head[u] = cnt1;
    e[++cnt1].to = u;e[cnt1].next = head[v];head[v] = cnt1;
}

void dfs1(ll u,ll faz,ll deep){
     dep[u] = deep;
     fa[u] = faz;
     siz[u] = 1;
     for(ll i = head[u];i;i=e[i].next){
        ll v = e[i].to;
        if(v != fa[u]){
            dfs1(v,u,deep+1);
            siz[u] += siz[v];
            if(son[u] == -1||siz[v] > siz[son[u]])
                son[u] = v;
        }
     }
}

void dfs2(ll u,ll t){
    top[u] = t;
    tid[u] = cnt;
    rk[cnt] = u;
    cnt++;
    if(son[u] == -1) return;
    dfs2(son[u],t);
    for(ll i = head[u];i;i = e[i].next){
        ll v = e[i].to;
        if(v != son[u]&&v != fa[u])
            dfs2(v,v);
    }
}

void pushup(ll rt){
    sum[rt].ls = sum[rt<<1].ls;
    sum[rt].rs = sum[rt<<1|1].rs;
    if(sum[rt<<1].rs==sum[rt<<1|1].ls) sum[rt].val = sum[rt<<1].val+sum[rt<<1|1].val-1;
    else sum[rt].val = sum[rt<<1].val + sum[rt<<1|1].val;
}

void build(ll l,ll r,ll rt){
    if(l == r){
        sum[rt].ls = a[rk[l]];
        sum[rt].rs = a[rk[l]];
        sum[rt].val = 1;
        ////cout<<l<<" "<<rk[l]<<" "<<a[rk[l]]<<endl;
        return ;
    }
    mid;
    build(lson);
    build(rson);
    pushup(rt);
}

void pushdown(ll rt){
    if(lazy[rt]){
    sum[rt<<1].ls = lazy[rt];
    sum[rt<<1|1].rs = lazy[rt];
    sum[rt<<1|1].ls = lazy[rt];
    sum[rt<<1].rs = lazy[rt];
    sum[rt<<1].val = 1;
    sum[rt<<1|1].val = 1;
    lazy[rt<<1] = lazy[rt<<1|1] = lazy[rt];
    lazy[rt] = 0;
    }
}

void update(ll L,ll R,ll c,ll l,ll r,ll rt){
     if(L <= l&&R >= r){
        sum[rt].val = 1;
        lazy[rt] = c;
        sum[rt].ls = c; sum[rt].rs = c;
        return ;
     }
     pushdown(rt);
     ll m = (l + r) >> 1;
     if(L <= m) update(L,R,c,lson);
     if(R > m)  update(L,R,c,rson);
     pushup(rt);
}

node1 query(ll L,ll R,ll l,ll r,ll rt){
    if(L <= l&&R >= r){
        return sum[rt];
    }
    pushdown(rt);
    ll m = (l + r) >> 1;
    if(L > m) return query(L,R,rson);
    if(R <= m) return query(L,R,lson);
    node1 t1 = query(L,m,lson);
    node1 t2 = query(m+1,R,rson);
    node1 t;
    t.ls = t1.ls;t.rs = t2.rs;
    if(t1.rs==t2.ls) t.val = t1.val+t2.val-1;
    else t.val = t1.val+t2.val;
    return t;
}

void cover(ll x,ll y,ll c){
    ll fx = top[x],fy = top[y];
    while(fx!=fy){
        if(dep[fx] < dep[fy]) swap(fx,fy),swap(x,y);
        update(tid[fx],tid[x],c,1,n,1);
        x = fa[fx];fx = top[x];
    }
    if(dep[x] < dep[y]) swap(x,y);
    update(tid[y],tid[x],c,1,n,1);
}

ll ask(ll x,ll y){
    ll sum = 0;
    ll lc = -1,rc=-1;
    ll fx = top[x],fy = top[y];
    node1 t;
    while(fx != fy){
        if(dep[fx] < dep[fy]){
            swap(fx,fy); swap(x,y); swap(lc,rc);
        }
        t = query(tid[fx],tid[x],1,n,1);
        sum += t.val - (lc==t.rs);
        x = fa[fx]; fx = top[x]; lc = t.ls;
    }
    if(dep[x] < dep[y]) swap(x,y),swap(lc,rc);
    t = query(tid[y],tid[x],1,n,1);
    sum += t.val - (lc==t.rs) - (rc==t.ls);  //当前是x-y区间与两端的区间相加，所以需要判两个
    return sum;
}


int main()
{
    ll u,v,x,y,m,z;
    memset(son,-1,sizeof(son));
    scanf("%lld%lld",&n,&m);
    cnt = 1;cnt1 = 1;
    for(ll i = 1;i <= n;i ++) {
        scanf("%lld",&x);
        a[i] = x+1;
    }
    for(ll i = 0;i < n-1;i++){
        scanf("%lld%lld",&u,&v);
        add(u,v);
    }
    dfs1(1,0,1); dfs2(1,1);
    build(1,n,1);
    char op[10];
    while(m--){
        scanf("%s",op);
        if(op[0] == ‘Q‘){
            scanf("%lld%lld",&x,&y);
            printf("%lld\n",ask(x,y));
        }
        else {
            scanf("%lld%lld%lld",&x,&y,&z);
            z++;
            cover(x,y,z);
        }
    }
    return 0;
}