2020 CCPC 长春 F. Strange Memory(树上启发式合并)

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)$，显然不行
考虑优化：计算节点 $u$ 的答案时，先计算其所有轻儿子的答案，不保留轻儿子的权值情况；
后计算其重儿子的答案，保留权值情况；
最后计算点u的答案时，则不需要再次遍历其重儿子所在子树
采用 DSU on Tree
优化后的复杂度为 $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;
}