P2486 [SDOI2011]染色 区间合并+树链剖分（加深对线段树的理解）

#include<bits/stdc++.h>
using namespace std;
const  int M=3e5+;
struct node{
    int l,r,cnt,lazy;
    node(int l1=,int r1=,int cnt1=,int lazy1=):l(l1),r(r1),cnt(cnt1),lazy(lazy1){}
}tree[M<<];
int fa[M],sz[M],deep[M],dfn[M],son[M],to[M],a[M],top[M],cnt,n;
char s[];
vector<int>g[M];
void dfs1(int u,int from){
    fa[u]=from;
    sz[u]=;
    deep[u]=deep[from]+;
    for(int i=;i<g[u].size();i++){

        int v=g[u][i];
        if(v!=from){
            dfs1(v,u);
            sz[u]+=sz[v];
            if(sz[v]>sz[son[u]])
                son[u]=v;
        }

    }
}
void dfs2(int u,int t){
    top[u]=t;
    dfn[u]=++cnt;
    to[cnt]=u;
    if(!son[u])
        return ;
    dfs2(son[u],t);
    for(int i=;i<g[u].size();i++){
        int v=g[u][i];
        if(v!=fa[u]&&v!=son[u])
            dfs2(v,v);
    }
}
void up(int root){
    tree[root].cnt=tree[root<<].cnt+tree[root<<|].cnt;
    if(tree[root<<].r==tree[root<<|].l)
        tree[root].cnt--;
    tree[root].l=tree[root<<].l;
    tree[root].r=tree[root<<|].r;
}
void build(int root,int l,int r){
    tree[root].lazy=;
    if(l==r){
        tree[root].l=tree[root].r=a[to[l]];
        tree[root].cnt=;
        return ;
    }
    int midd=(l+r)>>;
    build(root<<,l,midd);
    build(root<<|,midd+,r);
    up(root);
}
void pushdown(int root){
    tree[root<<]=tree[root<<|]=node(tree[root].l,tree[root].r,,tree[root].lazy);
    tree[root].lazy=;
}
void update(int L,int R,int x,int root,int l,int r){
    if(L<=l&&r<=R){
        tree[root]=node(x,x,,x);
        return ;
    }
    if(tree[root].lazy)
        pushdown(root);
    int midd=(l+r)>>;
    if(L<=midd)
        update(L,R,x,root<<,l,midd);
    if(R>midd)
        update(L,R,x,root<<|,midd+,r);
    up(root);
}
void add(int u,int v ,int w){
    int fu=top[u],fv=top[v];
    while(fu!=fv){
        if(deep[fu]>=deep[fv])
            update(dfn[fu],dfn[u],w,,,n),u=fa[fu],fu=top[u];
        else
            update(dfn[fv],dfn[v],w,,,n),v=fa[fv],fv=top[v];
    }
    if(dfn[u]<=dfn[v])
        update(dfn[u],dfn[v],w,,,n);
    else
        update(dfn[v],dfn[u],w,,,n);
}
node meger(node a,node b){
    if(!a.cnt)
        return b;
    if(!b.cnt)
        return a;
    node ans=node(,,,);
    ans.cnt=a.cnt+b.cnt;
    if(a.r==b.l)
        ans.cnt--;
    ans.l=a.l;
    ans.r=b.r;
    return ans;
}
node query(int L,int R,int root,int l,int r){
    if(L<=l&&r<=R){
        return tree[root];
    }
    if(tree[root].lazy)
        pushdown(root);
    int midd=(l+r)>>;
    node ans;
    if(L<=midd)
        ans=query(L,R,root<<,l,midd);
    if(R>midd)
        ans=meger(ans,query(L,R,root<<|,midd+,r));
    up(root);
    return ans;
}
int solve(int u,int v){
    node l,r;
    int fv=top[v],fu=top[u];
    while(fv!=fu){
        if(deep[fu]>=deep[fv])
            l=meger(query(dfn[fu],dfn[u],,,n),l),u=fa[fu],fu=top[u];
        else
            r=meger(query(dfn[fv],dfn[v],,,n),r),v=fa[fv],fv=top[v];
    }
    if(dfn[u]<=dfn[v])
        r=meger(query(dfn[u],dfn[v],,,n),r);
    else
        l=meger(query(dfn[v],dfn[u],,,n),l);
    swap(l.l,l.r);
    l=meger(l,r);
    return l.cnt;
}
int main(){
    int m;
    scanf("%d%d",&n,&m);
    for(int i=;i<=n;i++)
        scanf("%d",&a[i]);
    for(int i=;i<n;i++){
        int u,v;
        scanf("%d%d",&u,&v);
        g[u].push_back(v);
        g[v].push_back(u);
    }//cout<<"!!"<<endl;
    dfs1(,);
    dfs2(,);

    build(,,n);
    while(m--){
        int u,v,w;
        scanf("%s",s);
        if(s[]=='Q'){
            scanf("%d%d",&u,&v);
            printf("%d\n",solve(u,v));
        }
        else{
            scanf("%d%d%d",&u,&v,&w);
            add(u,v,w);
        }
    }
    return ;
}