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

权值统计

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

解析

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

  1. 以该点为路径的一个端点
  2. 路径经过该点。
    对于情况1,显然为\(dp_u\).
    对于情况2,可以发现,这样的路径的权值为(\(son_1\times son_2\times a_i\)).
    如何得到这个\(son_1\times son_2\)呢?
    我们发现:\((a_1+a_2+\cdots +a_n)^2 = \sum_{i=1}^n a_i^2 + 2\sum_{i=1}^n\sum_{j=i+1}^n a_i*a_j\).
    那么,求出\((\sum dp_v)^2\)和\(\sum dp_v^2\)后,相减并\(\div 2\),即是\(son_1\times 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;
}
上一篇:143. 最大异或对


下一篇:事件冒泡