宁波多校（四） B题 七彩魔法树（线段树）

#include<iostream>
#include<algorithm>
#include<stack>
#include<vector>
#include<cstring>
using namespace std;
typedef long long ll;
const int N=3e5+10;
int w[N],h[N],ne[N],e[N],idx;
int dfn[N],times,sz[N];
int id[N];
int n,m;
struct node{
    int l,r;
    ll cnt;
    ll lazy;
}tr[N<<2];
void add(int a,int b){
    e[idx]=b,ne[idx]=h[a],h[a]=idx++;
}
void dfs(int u,int fa){
    dfn[u]=++times;
    id[times]=u;
    sz[u]=1;
    for(int i=h[u];i!=-1;i=ne[i]){
        int j=e[i];
        if(j==fa)
            continue;
        dfs(j,u);
        sz[u]+=sz[j];
    }
}
void pushup(int u){
    tr[u].cnt=tr[u<<1].cnt|tr[u<<1|1].cnt;
}
void build(int u,int l,int r){
    if(l==r){
        tr[u]={l,r,1ll<<(w[id[l]]-1),0};
        return ;
    }
    else{
        tr[u]={l,r};
        int mid=l+r>>1;
        build(u<<1,l,mid);
        build(u<<1|1,mid+1,r);
        pushup(u);
    }
}
void pushdown(int u){
    tr[u<<1].cnt=tr[u].lazy;
    tr[u<<1|1].cnt=tr[u].lazy;
    tr[u<<1].lazy=tr[u].lazy;
    tr[u<<1|1].lazy=tr[u].lazy;
    tr[u].lazy=0;
}

void modify(int u,int l,int r,int c){
    if(tr[u].l>=l&&tr[u].r<=r){
        tr[u].cnt=(1ll<<c-1);
        tr[u].lazy=(1ll<<c-1);
        return ;
    }
    if(tr[u].lazy)
    pushdown(u);
    int mid=tr[u].l+tr[u].r>>1;
    if(l<=mid)
        modify(u<<1,l,r,c);
    if(r>mid)
        modify(u<<1|1,l,r,c);
    pushup(u);
}
ll query(int u,int l,int r){
    if(tr[u].l>=l&&tr[u].r<=r)
        return tr[u].cnt;
    if(tr[u].lazy)
    pushdown(u);
    int mid=tr[u].l+tr[u].r>>1;
    ll res=0;
    if(l<=mid)
        res=query(u<<1,l,r);
    if(r>mid)
        res|=query(u<<1|1,l,r);
    return res;

}
int check(ll x){
    int res=0;
    while(x){
        if(x&1)
            res++;
        x>>=1;
    }
    return res;
}
int main(){
    ios::sync_with_stdio(false);
    cin>>n>>m;
    int i;
    memset(h,-1,sizeof h);
    for(i=1;i<=n;i++){
        cin>>w[i];
    }
    for(i=1;i<n;i++){
        int u,v;
        cin>>u>>v;
        add(u,v);
        add(v,u);
    }
    dfs(1,-1);
    build(1,1,n);
    while(m--){
        int opt;
        cin>>opt;
        if(opt==1){
            int e,f;
            cin>>e>>f;
            modify(1,dfn[e],dfn[e]+sz[e]-1,f);
        }
        else{
            int q;
            cin>>q;
            cout<<check(query(1,dfn[q],dfn[q]+sz[q]-1))<<endl;
        }
    }
}