【YBTOJ】【树形dp】权值统计

权值统计

给出一个 (n) 个结点的无根树以及每个结点的权值，求出树的每一条路径的权值积的和，单独的一个结点也算作一条路径。

解析

要求每一条路径。
规定每条路径的权值为途径的点权之积。
对于这类问题，我们可以用 (dp_u) 表示以 (u) 为根的子树内所有答案。
但是，因为题目规定，权值为点权积。
考虑使用数学解决：
设(dp_u) 表示以(u)为根的子树内答案之和（以(u)为路径的一个端点）。
对于每个节点，可以发现：与该点有关的路径有两种情况：

1. 以该点为路径的一个端点
2. 路径经过该点。
   对于情况1，显然为 (dp_u).
   对于情况2，可以发现，这样的路径的权值为((son_1 imes son_2 imes a_i)).
   如何得到这个(son_1 imes son_2)呢？
   我们发现：((a_1+a_2+cdots +a_n)^2 = sumlimits_{i=1}^n a_i^2 + 2sumlimits_{i=1}^nsumlimits_{j=i+1}^n a_i*a_j).
   那么，求出((sum dp_v)^2)和(sum dp_v^2)后，相减并(div 2)，即是(son_1 imes son_2)。

代码

#include <bits/stdc++.h>
#define fo(a) freopen(a".in","r",stdin),freopen(a".out","w",stdout);
using namespace std;
const int INF = 0x3f3f3f3f,N = 1e5+5,mod = 10086;
typedef long long ll;
typedef unsigned long long ull;
inline ll read(){
	ll ret=0;char ch=' ',c=getchar();
	while(!(c>='0'&&c<='9'))ch=c,c=getchar();
	while(c>='0'&&c<='9')ret=(ret<<1)+(ret<<3)+c-'0',c=getchar();
	return ch=='-'?-ret:ret;
}
int n;
struct Edge{int to,nxt;}e[N<<1];
int head[N],ecnt = -1;
inline void add_edge(int u,int v){e[++ecnt] = (Edge){v,head[u]} , head[u] = ecnt;}

int a[N];
ll dp[N],ans;
void dfs(int u,int _f){
	ll sum=0,tot=0;
	for(int i = head[u] ; ~i ; i = e[i].nxt){
		int v = e[i].to;
		if(v == _f) continue;
		dfs(v,u);
		sum += dp[v],
		tot += dp[v] * dp[v];
	}
	dp[u] = ((sum + 1) * a[u]) % mod;
	(ans += dp[u] + (sum*sum-tot)/2*a[u]) %= mod;
}
signed main(){
	memset(head,-1,sizeof(head));
	n = read();
	for(int i = 1 ; i <= n ; i ++)
		a[i] = read();
	for(int i = 1 ; i < n ; i ++){
		int u = read() , v = read();
		add_edge(u,v);add_edge(v,u);
	}
	dfs(1,0);
	printf("%lld",ans);
	return 0;
}