Codeforces Round 899 (Div. 2)

D. Tree XOR

换根

#include<bits/stdc++.h>
using namespace std;
#define endl '\n'
#define yes cout << "Yes" << endl
#define no cout << "No" << endl
#define pii pair<int,int>
#define ll long long
#define pb push_back
#define ft first
#define se second
#define inf 0x3f3f3f3f
#define INF 0x3f3f3f3f3f3f3f3f

#define int long long
const int N = 200010;
int val[N];
int a[N];
int sz[N];
vector<int> G[N];
ll f[N][2];
int n;

ll ans[N];

void dfspre(int u, int fa){
sz[u] = 1;
    for(auto v:G[u]){
        if(v == fa) continue;
        dfspre(v, u);
        sz[u] += sz[v];
    }
}

void dfs(int u, int fa){
f[u][a[u]] = 0; f[u][!a[u]] = INF;
if(G[u].size() == 1) f[u][!a[u]] = 1;

    int s = 0;
    for(auto v:G[u]){
        if(v == fa) continue;
        dfs(v, u);
        f[u][a[u]] += f[v][a[u]];
    }
    f[u][!a[u]] = f[u][a[u]] + sz[u];
}
void dfs2(int u,int fa){
    for(auto v:G[u]){
        if(v == fa) continue;
        ans[v] += 1ll*ans[u] - 1ll*(ans[v] + sz[v]*(val[u]^val[v])) + 1ll*(n - sz[v]) * (val[u] ^ val[v]);
        dfs2(v, u);
    }
    
}

void solve(){
    cin >> n;
    for(int i = 1; i <= n; i ++) ans[i] = 0;
    for(int i = 1; i <= n; i ++) G[i].clear();

    for(int i = 1; i <= n; i ++) cin >> val[i];
    for(int i = 1; i <= n - 1;  i ++){
        int u, v; cin >> u >> v;
        G[u].pb(v);
        G[v].pb(u);
    }

  dfspre(1, -1);
    for(int j = 0; j < 20; j ++){
        for(int i = 1; i <= n; i ++){
            if(val[i] & (1 << j)) a[i] = 1;
            else a[i] = 0;
        }
        dfs(1, -1);
        // cout <<min(f[1][0], f[1][1])*(1<<j) << '\n';
        ans[1] += 1ll*min(f[1][0], f[1][1])*(1<<j);
    }
   dfs2(1, -1);
   for(int i = 1; i <= n; i ++){
    cout << ans[i] << ' ';
   }
    cout << '\n';
}

signed main(){
    std::ios::sync_with_stdio(false);
    int T = 1; 
    cin >> T;
    while(T--){
        solve();
    }
}