DSU on tree

存一下,可以用来对比找bug

#include<bits\stdc++.h>
using namespace std;
#define int long long
void in(int &x){
    int    y=1;char c=getchar();x=0;
    while(c<'0'||c>'9'){if(c=='-')y=-1;c=getchar();}
    while(c<='9'&&c>='0'){ x=(x<<1)+(x<<3)+c-'0';c=getchar();}
    x*=y;
}

const int _ = 1e6+10;
int n;int a[_];long long ans=0;
vector<int>g[_];
int Size[_],Mxsonid[_];
void dfs(int u,int fa){
    Size[u]=1;
    for(auto v:g[u]){
        if(v==fa)continue;
        dfs(v,u);
        Size[u]+=Size[v];
        if(Size[v]>Size[Mxsonid[u]])Mxsonid[u]=v;
    }
}
int cntt[_];
int cnt[_][20];
void push(int ax,int x){
    cntt[ax]++;
    for(int i=0;i<20;i++){
        if(x&(1<<i))cnt[ax][i]++;
    }
}
void pop(int ax,int x){
    cntt[ax]--;
    for(int i=0;i<20;i++){
        if(x&(1<<i))cnt[ax][i]--;
    }
}
void cal(int ax,int x){
    for(int i=0;i<20;i++) {
        if((x&(1<<i))){
            ans+=(1ll<<i)*(cntt[ax]-cnt[ax][i]);
        }else{
            ans+=(1ll<<i)*cnt[ax][i];
        }
    }
}
void add(int u,int fa);
void cal(int u,int fa,int val);
void del(int u,int fa);
void dfs2(int u,int fa,int flag){

    for(auto v:g[u]){
        if(v==Mxsonid[u]||v==fa)continue;
        dfs2(v,u,0);
    }
    if(Mxsonid[u])dfs2(Mxsonid[u],u,1);
    push(a[u],u);
    for(auto v:g[u]){
        if(v==Mxsonid[u]||v==fa)continue;
        cal(v,u,a[u]);
        add(v,u);
    }
    if(!flag)del(u,fa);
}


signed main(){
    in(n);
    for(int i=1;i<=n;i++)in(a[i]);
    for(int u,v,i=1;i<n;i++){
        in(u);in(v);
        g[u].push_back(v);g[v].push_back(u);
    }

    dfs(1,0);
    dfs2(1,0,1);
    cout<<ans;
    return 0;
}

void del(int u,int fa){
    pop(a[u],u);
    for(auto v:g[u]){
        if(v==fa)continue;
        del(v,u);
    }
}
void add(int u,int fa){
    push(a[u],u);
    for(auto v:g[u]){
        if(v==fa)continue;
        add(v,u);
    }
}
void cal(int u,int fa,int val){
    if((a[u]^val)<=1000000)cal(a[u]^val,u);
    for(auto v:g[u]){
        if(v==fa)continue;
        cal(v,u,val);
    }
}
上一篇:牛客练习赛91 C.魔法学院(hard version) (dsu区间染色)


下一篇:Dsu on tree 和 点分治