P2486 [SDOI2011]染色（树剖）区间覆盖+区间的连续段

#include<bits/stdc++.h>
using namespace std;

const int maxn = 100010;
vector<int>G[maxn];
int n , q ,cnt;
int siz[maxn] , wson[maxn],dep[maxn],fa[maxn],top[maxn],pos[maxn],ori[maxn];
void dfs1(int id , int F)
{
    siz[id]=1;
    for(int i=0 ; i<G[id].size() ; i++)
    {
        int v=G[id][i];
        if(v==F) continue;
        dep[v] = dep[id] + 1;
        fa[v] = id;
        dfs1(v,id);
        siz[id] += siz[v];
        if(siz[v] > siz[wson[id]]) wson[id] = v;
    }
}
void dfs2(int id , int TP)
{
    top[id] = TP;
    pos[id] = ++cnt;
    ori[cnt]=id;
    if(!wson[id]) return ;
    dfs2(wson[id],TP);
    for(int i=0 ; i<G[id].size() ; i++)
    {
        int v=G[id][i];
        if(v==fa[id] || v==wson[id]) continue;
        dfs2(v,v);
    }
}

int lc[maxn<<2] , rc[maxn<<2],col[maxn];
#define lson (id<<1)
#define rson (id<<1) | 1
struct no
{
    int l,r;
    int sum,c;///区间颜色总数 ， 叶子颜色
    int lazy;///儿子的颜色

}tree[maxn<<2];
void build(int id , int l , int r)
{
    tree[id].l=l;
    tree[id].r=r;
    if(l==r)
    {
        tree[id].c=col[ori[l]];//赋值：叶子颜色
        lc[id]=rc[id]=col[ori[l]];//赋值：区间左颜色和区间右颜色
        tree[id].sum=1;//颜色数为1
        return ;
    }
    int mid = (l+r)>>1;
    build(lson , l , mid);
    build(rson , mid+1 , r);
    tree[id].sum = tree[lson].sum + tree[rson].sum;
    if(rc[lson] == lc[rson]) tree[id].sum-=1;
    lc[id] = lc[lson];
    rc[id] = rc[rson];
}
void pushdown(int id)
{
    if(tree[id].lazy!=0 && tree[id].l != tree[id].r)
    {
        int c=tree[id].lazy;
        tree[lson].lazy = tree[rson].lazy=c;
        tree[lson].c = tree[rson].c = c;
        lc[lson]=lc[rson]=rc[lson]=rc[rson]=c;
        tree[lson].sum = tree[rson].sum=1;
        tree[id].lazy=0;
    }

}
void update(int id ,int c , int l , int r)
{
    pushdown(id);
    if(tree[id].l == l && tree[id].r==r)
    {
        tree[id].c=c;
        tree[id].lazy=c;
        tree[id].sum=1;
        lc[id]=rc[id]=c;
        return ;
    }
    int mid=tree[id].l + tree[id].r >> 1;
    if(mid < l)
        update(rson,c,l,r);
    else if(mid>=r)
        update(lson,c,l,r);
    else
    {
        update(lson,c,l,mid);
        update(rson,c,mid+1,r);
    }
    tree[id].sum=tree[lson].sum+tree[rson].sum;
    if(rc[lson]==lc[rson]) tree[id].sum-=1;
    lc[id]=lc[lson];
    rc[id]=rc[rson];
}
int query(int id , int l , int r)
{
    pushdown(id);
    if(tree[id].l == l && tree[id].r==r)
    {
        return tree[id].sum;
    }
    int mid=(tree[id].l + tree[id].r) >>1;
    if(mid < l)
        return query(rson,l,r);
    else if(mid >=r)
        return query(lson,l,r);
    else
    {
        int ret = query(lson,l,mid) + query(rson,mid+1,r);
        if(rc[lson]==lc[rson]) ret-=1;
        return ret;
    }
}
int Qc(int id , int l ,int r)
{
    pushdown(id);
    if(tree[id].l==l && tree[id].r==r)
    {
        return tree[id].c;
    }
    int mid=tree[id].l + tree[id].r >>1;
    if(mid < l) return Qc(rson,l,r);
    else return Qc(lson,l,r);
}
void uprange(int x , int y , int c)
{
    while(top[x]!=top[y])
    {
        if(dep[top[x]] < dep[top[y]]) swap(x,y);
        update(1,c,pos[top[x]],pos[x]);
        x=fa[top[x]];
    }
    if(dep[x] > dep[y]) swap(x,y);
    update(1,c,pos[x],pos[y]);
}
int Qsum(int x , int y)
{
    int ans=0,Cson,Cfa;
    while(top[x] != top[y])
    {

        if(dep[top[x]] < dep[top[y]]) swap(x,y);
        ans+=query(1,pos[top[x]],pos[x]);
        Cson=Qc(1,pos[top[x]],pos[top[x]]);
        Cfa=Qc(1,pos[fa[top[x]]] , pos[fa[top[x]]]);
        if(Cson==Cfa) ans-=1;
        x=fa[top[x]];
    }
    if(dep[x] > dep[y]) swap(x,y);
    ans+=query(1,pos[x] , pos[y]);
    return ans;
}
int main()
{
    scanf("%d%d",&n,&q);
    for(int i=1 ; i<=n ; i++) scanf("%d",&col[i]);
    int u,v;
    for(int i=1 ; i<n ; i++)
    {
        scanf("%d%d",&u,&v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs1(1,-1);
    dfs2(1,1);
    build(1,1,n);
    char ask;
    int c;
    while(q--)
    {
        cin>>ask;
        scanf("%d%d",&u,&v);
        if(ask==‘Q‘)
        {

            printf("%d\n",Qsum(u,v));
        }
        else
        {
            scanf("%d",&c);
            uprange(u,v,c);
        }
    }
}