[10.27_P2] 统计损失 (简单树形DP)

树形DP 简单题

---

Description

给定一棵树，每个节点有一个值。对于一条路径，它的值为路径上所有点的值的乘积。求出树上所有路径的值的和。
注意：单个点也算一条路径。

Input

第 1 行一个数 n，表示树的节点数。
第 2 行 n 个数，第 i 个数表示节点 i 的值。
第 3 到 n+1 行，每行两个数 u，v，表示 u，v 之间有一条边。

Output

包含一个数，表示答案模 10086 的值。

Sample Input

5
7 6 6 1 1
1 2
2 3
2 4
1 5

Sample Output

778

Hint

对于 100% 的数据，n <= 100000

---

题解

拿到题，我们可以很快看出，这应该是一道树形DP题。可以定义dp[i]为以节点 i 为根的子树的答案。然后从下往上更新。但是，等一等，如果直接这样更新，会有一点问题。如果这条路径不再往上走，或者只是经过节点 i ，向它的两个子树延伸，这样的数据向上传递后毫无作用，并且还会使答案错误。既然这样，我们可以在这个时候直接把这部分更新给ans，这就简单多了。

代码

#include <iostream>
#include <cstdio>
using namespace std;
#define adde(u,v) {e[cnt] = (edge){v,head[u]};head[u] = &e[cnt++];}


const int maxn = 1e5 + 5, mod = 10086;
int n,cnt;
int a[maxn];
int dp[maxn];
int ans = 0;

struct edge {
	int v;edge *next;
}e[maxn << 1], *head[maxn];

void dfs(int u,int fa) {
	dp[u] = a[u];
	ans = (ans + a[u]) % mod;
	for(edge *k = head[u];k;k = k->next)if(k->v != fa) {
		dfs(k->v,u);
		ans = (ans+dp[k->v]*dp[u]) % mod;
		dp[u] = (dp[u] + dp[k->v] * a[u]) % mod;
	}
}

int main() {
	scanf("%d",&n);
	for(int i = 1;i <= n;i++) scanf("%d",a + i),a[i] %= mod;
	for(int i = 1;i < n;i++) {
		int u,v;
		scanf("%d%d",&u,&v);adde(u,v);adde(v,u);
	}
	dfs(1,0);
	printf("%d",ans);
	
	return 0;
}