Descripton
Example
input
6
4 2 1 6 6 5
1 2
2 3
1 4
4 5
4 6
output
18
Solution
暴力计算统计答案复杂度为
O
(
n
2
)
O(n^2)
O(n2),显然不行
考虑优化:计算节点
u
u
u 的答案时,先计算其所有轻儿子的答案,不保留轻儿子的权值情况;
后计算其重儿子的答案,保留权值情况;
最后计算点u的答案时,则不需要再次遍历其重儿子所在子树
采用 DSU on Tree
优化后的复杂度为
O
(
n
l
o
g
n
)
O(nlogn)
O(nlogn)
Code
#include <bits/stdc++.h>
#define clr(a,b) memset(a,b,sizeof(a))
#define pb push_back
#define pii pair<int,int>
#define mpp make_pair
#define fi first
#define se second
typedef long long ll;
using namespace std;
const int maxn = 1e6 + 7;
ll ans;
int a[maxn];
vector<int>e[maxn];
int n,rt,sum,cnt;
int siz[maxn], maxp[maxn], son[maxn];
vector<int>judge[maxn];
pii tmp[maxn];
void pre_dfs(int u,int fa) {
siz[u] = 1;int maxp = 0;
for(auto v :e[u]) {
if(v == fa) continue;
pre_dfs(v,u);
siz[u] += siz[v];
if(siz[v] > maxp) {
maxp = siz[v];
son[u] = v;
}
}
}
void getv(int u,int fa) {
tmp[++cnt] = {u,a[u]};
for(auto v: e[u]) {
if(v == fa) continue;
getv(v,u);
}
}
queue<int>que;
void solve(int u, int fa,int flag) {
int d = a[u];
for(auto v : e[u]){
if(v == fa) continue;
cnt = 0;
if(v!=son[u]) getv(v,u);
else continue;
for(int j = 1;j <= cnt;++j) {
int x = d ^ tmp[j].se;
int sz = judge[x].size(); if(sz==0) continue;
for(int k = 0;k < sz;++k) {
ans += 1ll*(tmp[j].fi ^ judge[x][k]);
}
}
for(int j = 1;j <= cnt;++j) {
judge[tmp[j].se].pb(tmp[j].fi);que.push(tmp[j].se);
}
}
int sz = judge[0].size();
for(int k = 0;k < sz;++k) {
ans += 1ll * (u ^ judge[0][k]);
}
if(!flag) while(!que.empty()){judge[que.front()].clear();que.pop();}
else {
judge[a[u]].pb(u);que.push(a[u]);
}
}
void dfs(int u, int fa, int flag) {
for(auto v : e[u]) {
if(v == son[u] || v == fa) continue;
dfs(v,u,0);
}
if(son[u]) dfs(son[u],u,1);
solve(u,fa, flag);
}
int main() {
clr(son,0);
scanf("%d",&n);
for(int i = 1;i <= n;++i) scanf("%d",&a[i]);
for(int i = 1,u,v;i < n;++i) {
scanf("%d%d",&u,&v);
e[u].pb(v); e[v].pb(u);
}
pre_dfs(1,-1);
dfs(1,-1,0);
printf("%lld\n", ans);
return 0;
}