Codeforces 766E Mahmoud and a xor trip（树形DP）

题目链接 Mahmoud and a xor trip

树形DP。先考虑每个点到他本身的距离和，再算所有点两两距离和。

做的时候考虑二进制拆位即可。

#include <bits/stdc++.h>

using namespace std;

#define REP(i, n)     for(int i(0); i <  (n); ++i)
#define rep(i, a, b)  for(int i(a); i <= (b); ++i)
#define LL            long long

const int N = 100000 + 10;

int a[N];
int n, x, y;
LL dp[2][20][N];
LL ans;
vector <int> v[N];

void dfs(int x, int fa){
    REP(i, 20) dp[(a[x] >> i) & 1][i][x] = 1;

    for (auto it : v[x]){
        if (it == fa) continue;
        dfs(it, x);

        REP(i, 20) ans += (1LL << i) * (dp[0][i][x] * dp[1][i][it] + dp[1][i][x] * dp[0][i][it]);
        REP(i, 20){
            dp[((a[x] >> i) & 1)][i][x] += dp[0][i][it];
            dp[(((a[x] >> i) & 1)) ^ 1][i][x] += dp[1][i][it];
        }
    }
}

int main(){
    
    scanf("%d", &n);
    rep(i, 1, n) scanf("%d", a + i), ans += (LL)a[i];
    rep(i, 1, n - 1){
        scanf("%d%d", &x, &y);
        v[x].push_back(y); v[y].push_back(x);
    }

    dfs(1, 1);
    printf("%lld\n", ans);


    return 0;

}