F. Strange Memory(2022 CCPC 长春）

F. Strange Memory(2022 CCPC 长春）

tag： dsu on tree 位运算

题目链接

题意：

有一颗n个节点的树，其中1<=n<=10⁵ ,我们需要求解式子

即如果 a[i]^a[j] 和 i、j最近公共祖先的值相同 我们就加上 i^j

做法：

再做这题之前我并不会启发式合并，研究半天终于懂了个大概
在dsu on tree中我们将树的儿子节点们分为重儿子和轻儿子，重儿子的节点数最多且只有一个
现在我们可以将式子 a[i]^a[j] = a[lcm] 转化为 a[i]^a[lca] = a[j]，
对于一个以u为根的树，节点i属于重儿子，结点j属于轻儿子，那么显然lca(i,j) = u,我们可以遍历重儿子，再遍历轻儿子，o(n²)求贡献，但这显然不够优雅
我们可以通过位运算，设置数组cnt[a[i]^a[u]][21][2]，第一维表示 a[j] = a[i]^a[u]，将每个j通过二进制表示来算贡献
但这样只是算了重儿子对轻儿子的贡献，我们还要算轻儿子对轻儿子的贡献，只需要每次遍历完一个轻儿子后，将轻儿子和重儿子合并即可
为什么之后要删除轻儿子呢，因为此时的重儿子只是暂时的，如果它父亲不是重儿子，重儿子和轻儿子们又会被遍历一遍

代码：

#define fst std::ios::sync_with_stdio(false), std::cin.tie(0), std::cout << std::fixed << std::setprecision(20)
#define le "\n"
#define ll long long 
#include <bits/stdc++.h>
using namespace std;
const int N=1e5+50;
const int mod=998244353;
int a[N],son[N],sz[N];
vector<int> g[N];
int cnt[(1<<20)+7][21][2]; //为了统计和方便计算 a[j] = a[i]^a[lca] 的j，将每个j进行二进制表示来算贡献
ll ans;

void predfs(int u,int fa){
    sz[u] = 1;
    for(auto i: g[u]){
        if(i==fa) continue;
        predfs(i,u);
        sz[u] += sz[i];
        if(sz[i]>sz[son[u]]) son[u] = i;//计入重儿子
    }
}

void cal(int x,int y){
    for(int i=0;i<20;i++){
        if(y&(1<<i)) ans += (1ll<<i)*cnt[x][i][0];
        else ans += (1ll<<i)*cnt[x][i][1];
    }
}

void update(int u,int fa,int tf,int val){
    for(int i=0;i<20;i++){
        if(u&(1ll<<i)) cnt[a[u]][i][1] += val;
        else cnt[a[u]][i][0] += val;
    }
    for(auto v: g[u]){
        if(v==fa||v==son[tf]) continue;//注意这里是tf的重儿子而不是fa的重儿子
        update(v,u,tf,val);
    }
}

void count(int u,int fa,int tf){//tf标记待统计子树的根
    cal(a[u]^a[tf],u);
    for(auto v: g[u]){
        if(v==fa||v==son[tf]) continue;
        count(v,u,tf);//暴力遍历所有轻儿子的节点 
        if(u==tf) update(v,u,tf,1);//加上轻儿子,我们不仅要算轻儿子对重儿子的贡献，还要算轻儿子对重儿子的贡献
    }
    if(u==tf){
        for(auto v: g[u]){
            if(v==fa||v==son[tf]) continue;
            update(v,u,tf,-1);//删去轻儿子
        }
    }
}

void dfs(int u,int fa,int op){
    for(auto v: g[u]){
        if(v==son[u]||v==fa) continue;
        dfs(v,u,0);//首先遍历u的轻儿子算贡献
    }
    if(son[u]) dfs(son[u],u,1); //遍历u的重儿子算贡献
    count(u,fa,u);
    update(u,fa,u,1);//把u的二进制上传，此刻只保留了u和重儿子的二进制情况
    if(!op) update(u,fa,0,-1);//当前为u为fa的轻儿子，则删除贡献
}


int main() {
    fst;
    int n; cin>>n;
    for(int i=1;i<=n;i++) cin>>a[i];
    for(int i=1;i<n;i++){
        int u,v; cin>>u>>v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    predfs(1,0);
    dfs(1,0,1);
    cout<<ans<<le;
    return 0;
}