CF2241E Fair and Square

CF2241E Fair and Square

简单画图后发现路径上只有三个点的最近公共祖先被经过了三次，其余点都只被经过两次。

因此围绕这个最近公共祖先来统计，分为两种情况：

情况一，选这个点：那么就要在两个子树内各选一个点。

情况二，不选这个点：那么就要在三个子树内各选一个点。

在有根树中遍历，只需将父亲以上的部分也看成一棵子树就行。很容易理解这样是不重不漏的。

一开始我只想到了情况一的做法，这是一个很经典的套路，只需记录一个 \(sum\) 表示已加入的子树的大小之和，新的儿子 \(v\) 的贡献就是 \(sum \times siz_v\)。

那么情况二怎么做呢？其实是类似的，只需记录一个 \(las\) 表示选两个子树的贡献之和，新的儿子 \(v\) 的贡献就是 \(las \times siz_v\)。显然 \(las\) 就是情况一。

总结一下就是我们靠前缀和就可以完成贡献的统计，这个套路可以积累。

#include<bits/stdc++.h>
#define F(i,l,r) for(int i(l); i <= (r); ++ i)
#define G(i,r,l) for(int i(r); i >= (l); -- i)
using namespace std;
using ll = long long;
const int N = 3e5;
int siz[N], a[N];
vector<int> G[N], g[N]; 
int n;
ll ans; 
void dfs(int u, int fa){
	siz[u] = 1;
	for(auto v : G[u]){
		if(v == fa) continue;
		dfs(v, u);
		g[u].push_back(siz[v]);
		siz[u] += siz[v];
	}
	g[u].push_back(n - siz[u]);
}
void Main(){
	ans = 0;
	cin >> n;
	F(i, 1, n){
		cin >> a[i];
	} 
	F(i, 1, n - 1){
		int u, v;
		cin >> u >> v;
		G[u].push_back(v);
		G[v].push_back(u);
	}
	dfs(1, 0);
	F(u, 1, n){
		ll sum = 0, las = 0;
		int tmp = sqrt(a[u]);
		if(tmp * tmp < a[u]) continue; 
//		printf("%d\t", a[u]);
		if(g[u].size() >=2){
			for(auto x : g[u]){
				ans += sum * x;
				if(g[u].size() >= 3){
					ans += las * x;
					las += sum * x;
				}
				sum += x; 
			}
		}
	}
//	printf("\n");
	cout << ans << '\n';
	F(i, 1, n) {
		G[i].clear();
		g[i].clear();
	}
}
signed main(){
	ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
	int T;
	cin >> T;
	while(T --) Main();
	return fflush(0), 0;
}