树上启发式合并

L1-8 【十二奇策 - 07】子树

分数：20

题目描述

给定一棵拥有 \(n\) 个节点的树（根节点为 1 号节点），每个节点 \(i\) 都有一个权值 \(a_i\)。

对于每一个节点 \(i\)，定义其"贡献值" \(c_i\)，计算方式如下：
设以 \(i\) 为根的子树中，权值 \(x\) 出现的次数为 \(\text{cnt}_i(x)\)。设 \(S_i\) 为以 \(i\) 为根的子树中所有不同权值构成的集合，则 \(c_i\) 是子树中所有出现的权值 \(x\) 对应的 \((x \cdot \text{cnt}_i(x))\) 的按位异或和。也即：

\[c_i = \bigoplus_{x \in S_i} (x \cdot \text{cnt}_i(x)) \]

按位异或说明

对于两个非负整数 \(A\) 和 \(B\)，它们的按位异或 \(A \oplus B\) 定义如下：
在二进制表示下，考虑第 \(2^k\) 位 (\(k \ge 0\))。若 \(A\) 和 \(B\) 在该位上的数字恰好有一个为 1，则结果的该位为 1；否则为 0。
例如，\(3 \oplus 5 = 6\)，因为二进制 \(011 \oplus 101 = 110\)。

一般地，对于 \(k\) 个非负整数 \(p_1, p_2, \dots, p_k\)，它们的异或和定义为 \(((\dots((p_1 \oplus p_2) \oplus p_3) \oplus \dots) \oplus p_k)\)，且可以证明该结果与运算顺序无关。

请计算所有节点的 \(c_i\)（\(i=1,2,\dots,n\)）。

输入格式

第一行一个正整数 \(n(1 \le n \le 3 \times 10^5)\)，表示树上节点的个数。
接下来一行 \(n\) 个正整数，第 \(i\) 个正整数 \(a_i\) 表示编号为 \(i\) 的节点的权值 \((1 \le a_i \le n)\)。
接下来 \(n-1\) 行，每行两个正整数 \(u,v\)，表示节点 \(u\) 和节点 \(v\) 之间存在一条边。

保证输入数据构成一棵合法的树。

输出格式

输出一行 \(n\) 个非负整数，第 \(i\) 个数表示 \(c_i\) 的值。

输入样例

#include<bits/stdc++.h>
#define int long long
using namespace std;
const int N=3e5+5;
int a[N];
int head[N],nxt[N<<1],to[N<<1],tot;
inline void add(int x,int y){
    to[++tot]=y,nxt[tot]=head[x],head[x]=tot;
}
int cnt[N];
int ans[N];
int siz[N],son[N];
int Xor=0;
inline void dfs1(int now,int fa){
    siz[now]=1;
    son[now]=-1;
    int maxsz=0;
    for(int i=head[now];i;i=nxt[i]){
        int y=to[i];
        if(y==fa) continue;
        dfs1(y,now);
        siz[now]+=siz[y];
        if(siz[y]>maxsz){
            maxsz=siz[y];
            son[now]=y;
        }
    }
}
inline void update(int x,int delta){
    if(cnt[x]>0){
        Xor^=(x*cnt[x]);
    }
    cnt[x]+=delta;
    if(cnt[x]>0){
        Xor^=(x*cnt[x]);
    }
}
inline void dfsadd(int now,int fa){
    update(a[now],1);
    for(int i=head[now];i;i=nxt[i]){
        int y=to[i];
        if(y==fa) continue;
        dfsadd(y,now);
    }
}
inline void dfsclear(int now,int fa){
    update(a[now],-1);
    for(int i=head[now];i;i=nxt[i]){
        int y=to[i];
        if(y==fa) continue;
        dfsclear(y,now);
    }
}
inline void DFS(int now,int fa,int keep){
    for(int i=head[now];i;i=nxt[i]){
        int y=to[i];
        if(y==fa||y==son[now]) continue;
        DFS(y,now,false);
    }
    if(son[now]!=-1){
        DFS(son[now],now,true);
    }
    for(int i=head[now];i;i=nxt[i]){
        int y=to[i];
        if(y==fa||y==son[now]) continue;
        dfsadd(y,now);
    }
    update(a[now],1);
    ans[now]=Xor;
    if(!keep){
        dfsclear(now,fa);
        Xor=0;
    }
}
signed main(void){
    cin.tie(NULL)->sync_with_stdio(false);
    int n;
    cin>>n;
    for(int i=1;i<=n;i++){
        cin>>a[i];
    }
    for(int i=1;i<n;i++){
        int x,y;
        cin>>x>>y;
        add(x,y),add(y,x);
    }
    dfs1(1,0);
    DFS(1,0,false);
    for(int i=1;i<=n;i++){
        cout<<ans[i]<<" ";
    }
    cout<<'\n';
    
    return 0;
}

先树链刨分一样记录出来重儿子，然后跑 \(dfs\) 是先遍历轻儿子，回溯清空数据，再遍历重儿子，回溯不清空数据。每次操作是然后加上当前节点贡献，再加上所有轻儿子子树贡献。