【JZOJ3844】统计损失【树形dp】

题目大意：

题目链接：https://jzoj.net/senior/#main/show/3844
SJY有一天被LLT紧急召去计算一些可能的损失。LLT元首管理的SHB国的交通形成了一棵树，现在将会出现一颗陨石砸在SHB国中，并且陨石砸毁的必定是SHB国构成的交通树上的一条路径。SHB国的损失可表示为被砸毁的路径上的所有城市价值之积。现在还暂时无法确定陨石的掉落路线，所以LLT元首希望SJY能够告诉他SHB国在受到每一种砸毁方式后会受到的损失之和模10086之后的值。注意：单独一个节点也被认为是合法的路径。

---

思路：

设$f[x]$表示以$x$为根的子树内，所有节点到$i$的点权积之和。
那么枚举$x$的子节点$y$，有方程$f[x]=\sum f[y]\times a[x]$。
然后在每一次转移之前，$ans$都要加上$x$的其他子树与子树$y$之间每一条路径的权值积。所以就有$ans+=f[x]\times f[y]$
时间复杂度$O(n)$

---

代码：

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

const int N=100010,MOD=10086;
int n,ans,tot,f[N],a[N],head[N];

struct edge
{
	int next,to;
}e[N*2];

void add(int from,int to)
{
	e[++tot].to=to;
	e[tot].next=head[from];
	head[from]=tot;
}

void dfs(int x,int fa)
{
	f[x]=a[x]; ans+=a[x];
	for (int i=head[x];~i;i=e[i].next)
	{
		int v=e[i].to;
		if (v!=fa)
		{
			dfs(v,x);
			ans=(ans+f[v]*f[x])%MOD;
			f[x]=(f[x]+a[x]*f[v])%MOD;
		}
	}
}

int main()
{
	memset(head,-1,sizeof(head));
	scanf("%d",&n);
	for (int i=1;i<=n;i++)
		scanf("%d",&a[i]);
	for (int i=1,x,y;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		add(x,y); add(y,x);
	}
	dfs(1,0);
	printf("%lld",ans);
	return 0;
}