10.14TG T3 tree

这个题是树形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;
}