有大佬帮忙看看代码吗 有两个点T了 万分感谢QAQ
查看原帖
有大佬帮忙看看代码吗 有两个点T了 万分感谢QAQ
314588
cybercsc楼主2020/9/3 20:34
#include <bits/stdc++.h>
#define maxn 2000000+10
#define LL long long
using namespace std;
int MOD=998244353;
struct node{
    int w,to;
}e[2*maxn];
int ne[2*maxn],h[maxn],idx=0,ww[maxn],sz[maxn];
LL val[maxn],ans=0;
inline void add(int fa,int son,int w){
    e[idx].to=son;
    e[idx].w=w;
    ne[idx]=h[fa];
    h[fa]=idx;
    ++idx;
}
inline void dfs_fr(int fa,int pa){
    sz[fa]=1;//cout<<fa<<endl;
    for(register int i=h[fa];i!=-1;i=ne[i]){
        int v=e[i].to;
        if(v==pa)continue;
        dfs_fr(v,fa);sz[fa]+=sz[v];val[fa]+=val[v]+(LL)sz[v]*ww[fa];
    }
    val[fa]+=(LL)ww[fa];
    val[fa]%=MOD;
}
inline void dfs_se(int fa,int pa){
    for(register int i=h[fa];i!=-1;i=ne[i]){
        int v=e[i].to;
        if(v==pa)continue;
        ans=(LL)(ans+((LL)sz[v]*(val[fa]-val[v]-(LL)sz[v]*ww[fa]+(LL)MOD)%(MOD)+(LL)(sz[1]-sz[v])*val[v]%(MOD))*(LL)e[i].w%(MOD))%(MOD);
        val[v]=((LL)(sz[1]-sz[v])*ww[v]+val[fa]-(LL)sz[v]*ww[fa]+(LL)MOD)%MOD;
        //val[v]=(LL)(sz[1]-sz[v])*ww[v]+val[fa]-(LL)sz[v]*ww[fa];
        dfs_se(v,fa);
    }
}
int main()
{
    int n,x,y,z;cin>>n;memset(h,-1,sizeof(h));
    for(register int i=1;i<=n;i++)scanf("%d",&ww[i]);
    for(register int i=1;i<=n-1;i++){
        //cin>>x>>y>>z;
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);add(y,x,z);
    }
    dfs_fr(1,0);
    dfs_se(1,0);
    //cout<<val[1]<<" "<<val[2]<<" "<<val[3]<<endl;
    cout<<(ans*2)%MOD;
    return 0;
}

2020/9/3 20:34
加载中...