高数题 （树形DP？）

【问题描述】
π姐最近生了个漂亮宝宝休假在家，听lee姐介绍经验时发觉lee姐比高中反应迟钝了些，一问才知lee姐是因为连生了两儿子才这样的，lee姐的体会是一孕笨三年，π姐听了觉得闲着也是闲着，得找些挑战的事情做做，于是π姐最近开始刷起了高数题，她遇到了这样一道高数题。这道高数题里面有一棵N个点的树，树上每个点有点权，每条边有颜色。一条路径的权值是这条路径上所有点的点权和，一条合法的路径需要满足该路径上任意相邻的两条边颜色都不相同。问这棵树上所有合法路径的权值和是多少啊？（无向路径）
【输入格式】
第一行一个整数N，代表树上有多少个点。
接下来一行N个整数，代表树上每个点的权值。
接下来N-1行，每行三个整数S、E、C，代表S与E之间有一条颜色为C的边。
【输出格式】
一行一个整数，代表所求的值。
【输入样例】
6
6 2 3 7 1 4
1 2 1
1 3 2
1 4 3
2 5 1
2 6 2
【输出样例】
134
【样例解释】
1-2 Value: 8
1-3 Value: 9
1-4 Value:13
1-2-6 Value:12
2-1-3 Value:11
2-1-4 Value:15
2-5 Value:3
2-6 Value:6
3-1-4 Value:16
3-1-2-6 Value:15
4-1-2-6 Value:19
5-2-6 Value:7
【数据规模】
对与30%的数据，1≤N≤1000。
对于另外20%的数据，可用的颜色数不超过109且随机数据。
对于另外20%的数据，树的形态为一条链。
对于100%的数据，1≤N≤3*105，可用的颜色数不超过109，所有点权的大小不超过105。

题解：
这道题看上去很简单它也的确很简单
记录一个sum数组
$pa[i]表示以i为终点的路径总数$
$sum[i]表示以i为终点的所有路径的权值和$（从叶子往跟做，前提是相邻的边的颜色要不一样）
这个是处理能往上走的路径方案
然后不能往上走的也就是说以这个节点为中转站的
就是这个节点所有子节点
$pa[i]\ast sum[j]+pa[j]\ast sum[i]+pa[i]\ast pa[j]\ast w[x]$
$i,j,x$分别表示这个节点的任意两个子节点以及该子节点
因为每条新路径他们都要加一次，所以如上公式

结束！

#include<bits/stdc++.h>
using namespace std;
const int maxm=300010;
long long ans=0;
struct node
{
	long long nxt,v,color;
	long long w;
}e[maxm<<1];
long long cnt,head[maxm],w[maxm],pa[maxm],n;
long long sum[maxm],sump[maxm];
void add(int u,int v,int color)
{
	cnt++;
    e[cnt].v=v;
    e[cnt].nxt=head[u];
    e[cnt].color=color;
    head[u]=cnt;
}
void dfs(int x,int fa,int fr)
{
	for (int i=head[x]; i; i=e[i].nxt)
	  { 
	  	int y=e[i].v;
	  	if (y==fa) continue;
	  	dfs(y,x,e[i].color);
	  	if (fr!=e[i].color) 
		    {
			  pa[x]+=pa[y];
              sum[x]+=sum[y]+w[x]*pa[y];
			}else ans+=sum[y]+w[x]*pa[y];
	  }
	for (int i=head[x]; i; i=e[i].nxt)
	  for (int j=e[i].nxt; j; j=e[j].nxt)
	    if (i!=j&&e[i].color!=e[j].color&&e[i].v!=fa&&e[j].v!=fa)
	      {
	      	ans+=pa[e[i].v]*sum[e[j].v]+pa[e[j].v]*sum[e[i].v]+pa[e[i].v]*pa[e[j].v]*w[x];
	      }
	ans-=w[x];
	sum[x]+=w[x];
	pa[x]++;
}
int main()
{
	scanf("%d",&n);
	for (int i=1; i<=n; i++)
	    scanf("%d",&w[i]);
	for (int i=1; i<n; i++)
	  {
	  	int u,v,w;
	  	scanf("%d%d%d",&u,&v,&w);
	  	add(u,v,w);
	  	add(v,u,w);
	  }
	dfs(1,1,-1);
	for (int i=1; i<=n; i++)
	  ans+=sum[i];
	printf("%lld",ans);
}

注意一些细节，long long之类的

初测0，原因是第二个公式的有一个$j$打成了$i$，然后还过了样例没发现，还有是记录父亲连接爷爷的边的颜色
$e[i].color$打成了$i$，居然够了样例，然后时间没来的及造数据
难过qwq