C++(clang++ 11.0.1) 解法, 执行用时: 331ms, 内存消耗: 105540K, 提交时间: 2023-06-28 18:59:22

#include<bits/stdc++.h>
using namespace std;
const int N=200010;
int gi()
{
    char ch=getchar(); int f=1,x=0;
    while(ch<'0' || ch>'9'){if(ch=='-')f=-f; ch=getchar();}
    while(ch>='0' && ch<='9')x=x*10+ch-48,ch=getchar();
    return f*x;
}
int f[N][51],g[N][51];
int sz[N],c[N];
int n,head[N],cnt;
long long ans;
struct edge
{
    int to,next,h;
}e[N<<1];
void adde(int a,int b,int c)
{
    e[++cnt].to=b; e[cnt].next=head[a]; head[a]=cnt; e[cnt].h=c;
}
void dfs(int x,int fa)
{
    sz[x]=1; f[x][c[x]]=1;
    for(int i=head[x];i;i=e[i].next)
    {
        int v=e[i].to;
        if(v==fa)continue;
        dfs(v,x); sz[x]+=sz[v];
        for(int j=1;j<=50;j++)f[x][min(c[x],j)]+=f[v][j];
    }
}
void dp(int x,int fa)
{
    for(int i=head[x];i;i=e[i].next)
    {
        int v=e[i].to; 
        if(v==fa)continue;
        for(int j=1;j<=50;j++)g[v][min(c[x],j)]+=f[x][j]+g[x][j]-f[v][j];
        //for(int j=1;j<=50;j++)g[v][min(c[x],j)]-=f[v][j];
        dp(v,x);
        for(long long j=1;j<=50;j++)
            ans+=e[i].h*j*(1ll*sz[v]*g[v][j]+1ll*(n-sz[v])*f[v][j]);
    }
}
int main()
{
    n=gi(); 
    for(int i=1;i<=n;i++)c[i]=gi();
    int x,y,z;
    for(int i=1;i<n;i++)
    {
        x=gi(); y=gi(); z=gi(); 
        adde(x,y,z); adde(y,x,z);
    }
    dfs(1,0); dp(1,0);
    printf("%lld\n",ans);
    return 0;
}

C++(g++ 7.5.0) 解法, 执行用时: 509ms, 内存消耗: 79808K, 提交时间: 2023-06-23 20:35:38

#include<bits/stdc++.h>
using namespace std;
using ll=long long;
using ul=unsigned long long;
using pr=pair<int,int>;
mt19937_64 rg(random_device{}());
const int N=5e5+5,M=1e9+7;
int n,f[N][55],c[N],sz[N];
int hd[N],ed[N],to[N],w[N];
vector<int>lk[N];
ll ans;
void dp(int x,int y,int k){
    int i;sz[x]+=sz[y]*k;
    for(i=1;i<=50;++i)
        if(i>=c[x])f[x][c[x]]+=f[y][i]*k;
        else f[x][i]+=f[y][i]*k;
}
void dfs1(int x,int pr){
    int i,y;f[x][c[x]]=sz[x]=1;
    for(i=hd[x];i;i=to[i])
        if((y=ed[i])!=pr)
            dfs1(y,x),dp(x,y,1);
}
ll Sum(int x){
    ll res=0;
    for(int i=1;i<=50;++i)
        res+=i*f[x][i];
    return res;
}
void dfs2(int x,int pr){
    int i,y;
    for(i=hd[x];i;i=to[i])
        if((y=ed[i])!=pr){
            dp(x,y,-1),dp(y,x,1);
            dfs2(y,x),dp(y,x,-1);
            ans+=(Sum(x)*sz[y]+Sum(y)*sz[x])*w[i];
            dp(x,y,1);
        }
}
int main(){
    ios::sync_with_stdio(false);
    int i,j,k,l,r,x,y;
    for(i=1,cin>>n;i<=n;++i)cin>>c[i];
    for(i=1;i<n;++i){
        cin>>x>>y>>k;
        to[i+i]=hd[x],ed[hd[x]=i+i]=y;
        to[i+i+1]=hd[y],ed[hd[y]=i+i+1]=x;
        w[i+i]=w[i+i+1]=k;
    }dfs1(1,0),dfs2(1,0);
    printf("%lld\n",ans);
    return 0;
}