热狗树 树形dp（中国石油大学我要变强第九场）

题目描述
“我是番茄酱！”
“我是黄芥末酱！”
“合在一起就是——美式热狗上加的，那个！“
热狗树上的每个节点都涂有番茄酱或者黄芥末酱中的一种，这样热狗树就变得美味了~LiMn2O4构造了一颗热狗树，他想知道这棵树的美味程度。
一个热狗树的美味程度，定义为每个节点到其他和自己口味不一样的节点的最短距离之和的和，也就是任意两个口味不同的节点之间的路径长度和。请你求出这颗树的美味值，并且答案对998244353取模。

输入
第一行一个正整数N。
第二行有N个数，他们的值在[0,1]范围内取，其中1代表是涂了番茄酱，0代表是涂了黄芥末酱。
接下来N-1行，每一行有三个数a,b,w。代表节点a到b有边，路径的权值是w。输入数据保证是一棵树。
N≤100000,1≤a,b≤N,w<1 000 000 000

输出
输出答案对998244353取模后的结果。

样例输入
复制样例数据
3
1 0 1
1 2 1
1 3 2
样例输出
8

提示
树（tree）是包含n（n>=0）个结点的有穷集，其中：
（1）每个元素称为结点（node）；
（2）有一个特定的结点被称为根结点或树根（root）。
（3）除根结点之外的其余数据元素被分为m（m≥0）个互不相交的集合T1，T2，……Tm-1，其中每一个集合Ti（1<=i<=m）本身也是一棵树，被称作原树的子树（subtree）。
补一下以前的坑，当时比赛的时候没有做到这个题。。今天想起来了，就来补一下。
统计不同颜色的节点互相到达的距离和。要是我们去枚举点到点的距离，这样肯定是不行的，因为那样是O（n^2） 的时间复杂度。一般这样我么就去枚举边，因为边的条数不多。那么我们去考虑一下每条边的贡献。每一条边的贡献就等于这个节点的子树中0的个数*（1的个数-子树中1的个数）* w+子树中1的个数*（0的个数-子树中0的个数）。
代码如下：

#include<bits/stdc++.h>
#define ll long long
#define mod 998244353
using namespace std;

const int maxx=1e5+100;
struct edge{
	int to;
	int next;
	ll w;
}e[maxx<<1];
int head[maxx],col[maxx];
ll dp[maxx][2],ans,cnt1,cnt0;
int n,tot;

void add(int u,int v,ll w)
{
	e[tot].to=v;e[tot].w=w;e[tot].next=head[u];head[u]=tot++;
}

inline void dfs(int u,int f)
{
	dp[u][col[u]]=1;
	for(int i=head[u];i!=-1;i=e[i].next)
	{
		int to=e[i].to;
		if(to==f) continue;
		dfs(to,u);
		ans=(ans+(dp[to][1]*(cnt0-dp[to][0])%mod)*e[i].w%mod)%mod;
		ans=(ans+(dp[to][0]*(cnt1-dp[to][1])%mod)*e[i].w%mod)%mod;
		dp[u][1]+=dp[to][1];//记录1的个数
		dp[u][0]+=dp[to][0];//记录0的个数
	}
}
int main()
{
	int l,r;ll w;
	while(~scanf("%d",&n))
	{
		tot=cnt1=cnt0=0;
		memset(head,-1,sizeof(head));
		for(int i=1;i<=n;i++) scanf("%d",&col[i]),cnt1+=col[i];
		cnt0=n-cnt1;
		for(int i=1;i<n;i++)
		{
			scanf("%d%d%lld",&l,&r,&w);
			add(l,r,w);
			add(r,l,w);
		}
		ans=0;
		dfs(1,0);
		printf("%lld\n",ans*2%mod);//每一条边的贡献我们在dfs的时候只求了一遍，但是实际上一条边该是贡献了两遍价值
	}
	return 0;
}

努力加油a啊，(^o)/~