BZOJ 2243: [SDOI2011]染色 (树链剖分，点权，线段树)

题目链接

树链剖分，然后直接跑一遍，线段树，本题是记录区间的段数，所以线段树节点存下每个区间的左右端点的值，在设一个color懒惰标志，-1表示没有涂色即可

//#pragma comment(linker, "/STACK:102400000,102400000")
#include<bits/stdc++.h>
using namespace std;
#define cl(a,b) memset(a,b,sizeof(a))
#define fastIO ios::sync_with_stdio(false);cin.tie(0);
#define ll  long long
#define pb push_back
#define gcd __gcd


#define For(i,j,k) for(int i=(j);i<=k;i++)
#define lowbit(i) (i&(-i))
#define _(x) printf("%d\n",x)

typedef vector<ll> vec;
typedef pair<int,int> PI;
const double EPS = 1e-8;
const int maxn = 2e5+100;
const int inf  = 1 << 28;

int n,m;
int a[maxn];
struct Edge{
    int to,next;
}es[maxn<<2];
int head[maxn],tot;
void addEdge(int u,int v){
    es[tot].to=v;es[tot].next=head[u];head[u]=tot++;
}

int num[maxn],top[maxn],son[maxn],deep[maxn],p[maxn],fp[maxn],fa[maxn],pos;
void dfs(int u,int pre,int d){
    num[u]=1;
    fa[u]=pre;
    deep[u]=d;
    for(int i=head[u];~i;i=es[i].next){
        int v=es[i].to;
        if(v!=pre){
            dfs(v,u,d+1);
            num[u]+=num[v];
            if(son[u]==-1||num[son[u]]<num[v])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;i=es[i].next){
        int v=es[i].to;
        if(v!=fa[u]&&v!=son[u])getpos(v,v);
    }
}

void init(){
    tot=0;
    cl(head,-1);
    pos=1;
    cl(son,-1);
}

struct node{
    int sum,color,lc,rc;
}seg[maxn<<2];

void pushup(int rt){
    seg[rt].sum=seg[rt<<1].sum+seg[rt<<1|1].sum-(seg[rt<<1].rc==seg[rt<<1|1].lc);
    seg[rt].lc=seg[rt<<1].lc;
    seg[rt].rc=seg[rt<<1|1].rc;
}
void pushdown(int rt){
    if(seg[rt].color!=-1){
        seg[rt<<1].color=seg[rt<<1|1].color=seg[rt].color;
        seg[rt<<1].sum=seg[rt<<1|1].sum=1;
        seg[rt<<1].lc=seg[rt<<1].rc=seg[rt].color;
        seg[rt<<1|1].lc=seg[rt<<1|1].rc=seg[rt].color;
        seg[rt].color=-1;
    }
}
void build(int rt,int l,int r){
    seg[rt].color=-1;
    if(l==r){
        seg[rt].sum=1;
        seg[rt].color=seg[rt].lc=seg[rt].rc=a[fp[l]];
        return;
    }
    int mid=l+r>>1;
    build(rt<<1,l,mid);
    build(rt<<1|1,mid+1,r);
    pushup(rt);
}
void update(int rt,int l,int r,int x,int y,int val){
    if(x<=l&&r<=y){
        seg[rt].color=seg[rt].rc=seg[rt].lc=val;
        seg[rt].sum=1;return;
    }
    pushdown(rt);
    int mid=l+r>>1;
    if(x<=mid) update(rt<<1,l,mid,x,y,val);
    if(y>mid) update(rt<<1|1,mid+1,r,x,y,val);
    pushup(rt);
}

int query(int rt,int l,int r,int x,int y){
    if(x<=l&&r<=y)return seg[rt].sum;
    pushdown(rt);
    int mid=l+r>>1;
    int ans=0;
    if(x<=mid)ans+=query(rt<<1,l,mid,x,y);
    if(y>mid) ans+=query(rt<<1|1,mid+1,r,x,y);
    if(x<=mid&&y>mid && seg[rt<<1].rc==seg[rt<<1|1].lc)--ans;
    return ans;
}

int getcolor(int rt,int l,int r,int k){
    if(l==r&&k==l)return seg[rt].color;
    pushdown(rt);
    int mid=l+r>>1;
    if(k<=mid) return getcolor(rt<<1,l,mid,k);
    else return getcolor(rt<<1|1,mid+1,r,k);
}

void modify(int u,int v,int x){
    while(top[u]!=top[v]){
        if(deep[top[u]]<deep[top[v]])swap(u,v);
        update(1,1,pos-1,p[top[u]],p[u],x);
        u=fa[top[u]];
    }
    if(deep[u]>deep[v])swap(u,v);
    update(1,1,pos-1,p[u],p[v],x);
}
int findSum(int u,int v){
    int ans=0;
    while(top[u]!=top[v]){
        if(deep[top[u]]<deep[top[v]])swap(u,v);
        ans+=query(1,1,pos-1,p[top[u]],p[u]);
        if(getcolor(1,1,pos-1,p[fa[top[u]]]) == getcolor(1,1,pos-1,p[top[u]]))ans--;
        u=fa[top[u]];
    }
    if(deep[u]>deep[v])swap(u,v);
    return ans+query(1,1,pos-1,p[u],p[v]);
}

int main(){
    //freopen("H:\\paint.in","r",stdin);
    //freopen("H:\\cc.out","w",stdout);
    while(~scanf("%d%d",&n,&m)){
        init();
        for(int i=1;i<=n;i++){
            scanf("%d",&a[i]);
        }
        for(int i=0;i<n;i++){
            int x,y;scanf("%d%d",&x,&y);
            addEdge(x,y);addEdge(y,x);
        }
        dfs(1,0,0);getpos(1,1);
        build(1,1,pos-1);
        while(m--){
            char op[3];
            scanf("%s",op);
            if(op[0]=='C'){
                int x,y,z;scanf("%d%d%d",&x,&y,&z);
                modify(x,y,z);
            }
            else {
                int u,v;scanf("%d%d",&u,&v);
                printf("%d\n",findSum(u,v));
            }
        }
    }
    return 0;
}

/*

6 5
2 2 1 2 1 1
1 2
1 3
2 4
2 5
2 6
Q 3 5
C 2 1 1
Q 3 5
C 5 1 2
Q 3 5

*/