题目描述
给定一棵 $n$ 个点的树,点 $i$ 的点权为 $a_i$。定义 $f(u, v)$ 为从树链 $u, v$ 上所有点的点权的异或和。
求 $\sum\limits_{u=1}^n\sum\limits_{v=u}^n f(u, v)$
限制:
- $1 \leqslant n \leqslant 10^5$
- $0 \leqslant a_i \leqslant 10^6$
算法
(树型dp) $O(n\log \max(a_i))$
本题实际上是 异或和之和 的树上版本
考虑树形dp
记 dp[v][b][x]
表示在以点 $v$ 为根的子树中从点 $v$ 出发且二进制下第 $b$ 位的异或和为 $x$ 的树链个数
C++ 代码
#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); ++i)
using namespace std;
using ll = long long;
int main() {
cin.tie(nullptr) -> sync_with_stdio(false);
int n;
cin >> n;
vector<int> a(n);
rep(i, n) cin >> a[i];
vector<vector<int>> to(n);
rep(i, n-1) {
int a, b;
cin >> a >> b;
--a; --b;
to[a].push_back(b);
to[b].push_back(a);
}
ll ans = 0;
vector dp(n, vector<ll>(2));
auto dfs = [&](auto& f, int v, int p, int b) -> void {
int x = a[v]>>b&1;
dp[v][0] = dp[v][1] = 0;
dp[v][x] = 1;
ans += dp[v][1]*(1<<b);
for (int u : to[v]) {
if (u == p) continue;
f(f, u, v, b);
ll now = 0;
now += dp[v][0]*dp[u][1];
now += dp[v][1]*dp[u][0];
ans += now*(1<<b);
dp[v][0] += dp[u][x];
dp[v][1] += dp[u][!x];
}
};
rep(i, 20) dfs(dfs, 0, -1, i);
cout << ans << '\n';
return 0;
}