这个题是树形dp
我们可以记录两个量,sum和cnt
sum(x)代表x的子树中的合法路径长度总和,cnt(x)代表x的子树中的合法路径数量。
那么就好dp了:
#include<cstdio>
#include<cstring>
#include<iostream>
#include<cmath>
#include<algorithm>
#include<cstdlib>
#define ll long long
using namespace std;
inline int read(){
int x=0;char ch=' ';int f=1;
while(ch!='-'&&(ch<'0'||ch>'9'))ch=getchar();
if(ch=='-')f=-1,ch=getchar();
while(ch>='0'&&ch<='9')x=x*10+ch-'0',ch=getchar();
return x*f;
}
const int N=300001;
struct edge{
int to,next,col;
}e[N*2];
int n,tot;
int head[N],v[N];
inline void addedge(int x,int y,int c){
e[++tot].to=y;e[tot].next=head[x];e[tot].col=c;head[x]=tot;
}
ll sum[N],cnt[N],ans;
inline void dfs(int x,int fa,int c){
cnt[x]=1;sum[x]=v[x];
for(int i=head[x];i;i=e[i].next){
int u=e[i].to;
if(u==fa)continue;
dfs(u,x,e[i].col);
if(c!=e[i].col){
sum[x]+=sum[u]+cnt[u]*v[x];
cnt[x]+=cnt[u];
}
ans+=cnt[u]*v[x]+sum[u];
}
for(int i=head[x];i;i=e[i].next){
int u=e[i].to;
if(u==fa)continue;
for(int j=i;j;j=e[j].next){
int h=e[j].to;
if(h==fa)continue;
if(e[i].col!=e[j].col){
ans+=sum[u]*cnt[h]+sum[h]*cnt[u]+v[x]*cnt[u]*cnt[h];
}
}
}
}
int main(){
n=read();
for(int i=1;i<=n;++i){
v[i]=read();
}
for(int i=1;i<n;i++){
int x=read(),y=read(),c=read();
addedge(x,y,c);addedge(y,x,c);
}
dfs(1,0,0);
printf("%lld",ans);
return 0;
}