CF gym102832 F. Strange Memory(dsu on tree+二进制拆分)

https://vjudge.net/problem/Gym-102832F

题意：给定n个结点，每个结点有一个权值a_i，求下列值。

思路：
因为a_i≠0，所以ai⊕aj≠ai或者aj，所以符合条件的i，j对的lca(i,j)不会是它们其中一个。当以一个点作为lca时，能产生贡献的i和j只能来自这个点的不同儿子子树。
考虑以当前结点作为lca时，不同儿子子树对答案产生的贡献：我们需要记录这棵子树上每种权值 有 对应哪些下标值，因为是计算异或值和，这里可以进行二进制拆分，只记录权值相同时每一位上下标值为0或1的下标值的个数，这样就可以分别计算每一位的贡献。计算完一棵儿子子树的贡献后，再将它的信息合并到cnt数组中。
然后利用dsu on tree优化：当当前结点是它父亲的重儿子时，可以不用清空cnt数组。

注意：
①算贡献和记录下标值信息时，前面的0位也都要考虑。
②两个1e6级别的数 异或结果可能是2e6级别的

#include<algorithm>
#include<iostream>
using namespace std;
#define ll long long
#define ull unsigned long long
#define PII pair<int,int>
#define mid ((l + r)>>1)
#define chl (root<<1)
#define chr (root<<1|1)
#define lowbit(x) ( x&(-x) )
const int manx = 1e5 + 10;
const int manx2 = 1e6 + 10;
const int INF = 2e9;
const int mod = 1e4+7;

int num[manx],hson[manx],sz[manx];
int cnt[manx2][18][2];
//cnt[i][j][]表示权值为i的这些点中 第j为1、0的下标值个数
ll ans;
int cou,head[manx];
struct node
{
    int e,bf;
}edge[manx<<2];
void init(int n)
{
    cou=0,ans=0;
    for(int i=1;i<=n;i++){
        head[i]=-1;
        hson[i]=0;
        sz[i]=1;
    }
    sz[0]=-1;
}
void add1(int s,int e)
{
    edge[cou]=node{e,head[s]};
    head[s]=cou++;
}
void dfs_list(int s,int fa)//找每个结点重儿子
{
    for(int i=head[s];~i;i=edge[i].bf){
        int e=edge[i].e;
        if(e==fa)continue;
        dfs_list(e,s);
        sz[s]+=sz[e];
        if(sz[e]>sz[hson[s]])
            hson[s]=e;
    }
}
void add(int s,int fa,int root_val)
{
    int val=num[s]^root_val;
    if(val<=1e6)
    for(int i=1,j=1;j<=17;i<<=1,j++){
        if(s&i)ans+=1ll*i*cnt[val][j][0];
        else ans+=1ll*i*cnt[val][j][1];
    }
    for(int i=head[s];~i;i=edge[i].bf){
        int e=edge[i].e;
        if(e==fa)continue;
        add(e,s,root_val);
    }
}
void update(int s,int fa,int k)
{
    for(int i=1,j=1;j<=17;i<<=1,j++){
        if(i&s)cnt[num[s]][j][1]+=k;
        else cnt[num[s]][j][0]+=k;
    }
    for(int i=head[s];~i;i=edge[i].bf){
        int e=edge[i].e;
        if(e==fa)continue;
        update(e,s,k);
    }
}
void dfs_Dot(int s,int fa,int keep)
{
    for(int i=head[s];~i;i=edge[i].bf){
        int e=edge[i].e;
        if(e!=fa&&e!=hson[s])
            dfs_Dot(e,s,0);//1、遍历重儿子，不记录cnt
    }
    if(hson[s])dfs_Dot(hson[s],s,1);//2、遍历重儿子，记录cnt
    //3、再遍历一遍轻儿子，找最近公共祖先为结点s为的子答案，记录cnt
    for(int i=head[s];~i;i=edge[i].bf){
        int e=edge[i].e;
        if(e!=fa&&e!=hson[s]){
            add(e,s,num[s]);
            update(e,s,1);
        }
    }
    for(int i=1,j=1;j<=17;i<<=1,j++){
        if(i&s)cnt[num[s]][j][1]++;
        else cnt[num[s]][j][0]++;
    }
    if(!keep)update(s,fa,-1);//当前结点不是重儿子(keep==0)时，清空cnt
}
int main()
{
    int n,s,e;
    scanf("%d",&n);
    init(n);
    for(int i=1;i<=n;i++)
        scanf("%d",&num[i]);
    for(int i=1;i<n;i++){
        scanf("%d%d",&s,&e);
        add1(s,e);
        add1(e,s);
    }
    dfs_list(1,0);
    dfs_Dot(1,0,1);
    printf("%lld\n",ans);
    return 0;
}