WA*2 求调,悬关 3 个
查看原帖
WA*2 求调,悬关 3 个
706523
AlicX楼主2024/6/3 19:23

rt.

#include<bits/stdc++.h> 
#define int long long  
#define x first 
#define y second 
#define il inline 
#define debug() puts("-----") 
using namespace std; 
typedef pair<int,int> pii; 
il int read(){ 
	int x=0,f=1; char ch=getchar(); 
	while(ch<'0'||ch>'9'){ if(ch=='-') f=-1; ch=getchar(); } 
	while(ch>='0'&&ch<='9') x=(x<<1)+(x<<3)+(ch^48),ch=getchar(); 
	return x*f; 
} 
const int N=2e6+10; 
const int Mod=998244353; 
int n; 
int a[N]; 
int f[N]; 
int siz[N]; 
int h[N],idx=0; 
struct Edge{ 
	int w; 
	int to,ne; 
}e[N<<1]; 
il void add(int u,int v,int w){ 
	e[idx].w=w,e[idx].to=v,e[idx].ne=h[u],h[u]=idx++; 
} 
int sv[N],sw[N]; 
int gv[N],gw[N]; 
il void dfs(int u,int fa,int s1,int s2){ 
	sv[u]=a[u]; 
	f[1]=(f[1]+s1*s2%Mod)%Mod; 
	siz[u]=1; gv[1]+=s1,gw[1]+=s2; 
	for(int i=h[u];i!=-1;i=e[i].ne){ 
		int to=e[i].to; 
		if(to==fa) continue; 
		dfs(to,u,s1+a[to],s2+e[i].w); 
		siz[u]+=siz[to]; 
		sv[u]+=sv[to]+siz[to]*a[u]; 
		sw[u]+=sw[to]+siz[to]*e[i].w; 
	} 
} 
il void calc(int u,int fa,int w){ 
	if(u!=1){ 
		f[u]=f[fa]; 

		gw[u]=gw[fa]+(n-siz[u])*w-siz[u]*w; 
		gv[u]=gv[fa]+(n-siz[u])*a[u]-siz[u]*a[fa]; 
		// 子树外 
		f[u]=(f[u]+(n-siz[u])*w*a[u]%Mod)%Mod; 
		f[u]=(f[u]+(gw[fa]-sw[u]-w*siz[u])*a[u]%Mod+(gv[fa]-sv[u]-siz[u]*a[fa])*w%Mod)%Mod; 
		// 子树内 
		f[u]=(f[u]+siz[u]*w*a[fa]%Mod)%Mod; 
		f[u]=(f[u]-(sw[u]+w*siz[u])*a[fa]%Mod-(sv[u]+siz[u]*a[fa])*w%Mod)%Mod; f[u]=(f[u]%Mod+Mod)%Mod; 
	} for(int i=h[u];i!=-1;i=e[i].ne){ 
		int to=e[i].to; 
		if(to==fa) continue; 
		calc(to,u,e[i].w); 
	} 
} 
signed main(){ 
	memset(h,-1,sizeof h); 
	n=read(); int ans=0; 
	for(int i=1;i<=n;i++) a[i]=read(); 
	for(int i=1,u,v,w;i<n;i++) u=read(),v=read(),w=read(),add(u,v,w),add(v,u,w); 
	dfs(1,0,a[1],0); calc(1,0,0); 
	for(int i=1;i<=n;i++) ans=(ans+f[i])%Mod; 
	printf("%lld\n",(ans+Mod)%Mod); return 0; 
} 
2024/6/3 19:23
加载中...