T了两个点,求卡常
查看原帖
T了两个点,求卡常
812955
System_Error_/se楼主2024/3/25 21:32
#include<bits/stdc++.h>
#define int long long
using namespace std;
const int N=2e6+5,p=998244353;
int n,v[N],sumv[N],sumw[N],sz[N],ans[N],anss[N],sum=0;
bool vis[N];
struct node {
	int x,y;
	node(int X=0,int Y=0) {
		x=X,y=Y;
	}
};
vector<node> e[N];
inline void calc(int x,int pre) {
	sz[x]=1;
	for(int i=0; i<e[x].size(); i++) {
		int xx=e[x][i].x,yy=e[x][i].y;
		if(xx==pre) continue;
		calc(xx,x);
		sumv[x]=(sumv[x]+sumv[xx])%p;
		sumw[x]=(sumw[x]+yy*sz[xx]%p+sumw[xx])%p;
		sz[x]+=sz[xx];
		ans[x]=(ans[x]+ans[xx]+v[x]*sumw[xx]%p+yy*((sumv[xx]-v[xx]+p)%p)%p+v[x]*yy%p*((sz[xx]-1+p)%p)%p+(v[x]+v[xx])*yy%p)%p;
	}
	sumv[x]=(sumv[x]+sz[x]*v[x]%p)%p;
}
inline void change(int x,int y,int w) {
	ans[x]=(ans[x]-ans[y]-v[x]*sumw[y]%p-w*((sumv[y]-v[y]+p)%p)%p-v[x]*w%p*((sz[y]-1+p)%p)%p-(v[x]+v[y])*w%p+5*p)%p;
	sumv[x]=(sumv[x]-sz[y]*v[x]%p-sumv[y]+p+p)%p;
	sumv[y]=(sumv[y]+(n-sz[y])*v[y]%p+sumv[x])%p;
	sumw[x]=(sumw[x]-sz[y]*w%p-sumw[y]+p+p)%p;
	sumw[y]=(sumw[y]+(n-sz[y])*w%p+sumw[x])%p;
	ans[y]=(ans[y]+ans[x]+v[y]*sumw[x]%p+w*((sumv[x]-v[x]+p)%p)%p+v[y]*w%p*((n-sz[y]-1+p+p)%p)%p+(v[x]+v[y])*w%p)%p;
	sz[x]=n-sz[y],sz[y]=n;
}
inline void dfs(int x) {
	vis[x]=true,anss[x]=ans[x];
	for(int i=0; i<e[x].size(); i++) {
		int xx=e[x][i].x,yy=e[x][i].y;
		if(vis[xx]) continue;
		int aa=ans[x],ab=ans[xx],ac=sumv[x],ad=sumv[xx],ae=sumw[x],af=sumw[xx],ag=sz[x],ah=sz[xx];
		change(x,xx,yy);
		dfs(xx);
		ans[x]=aa,ans[xx]=ab,sumv[x]=ac,sumv[xx]=ad,sumw[x]=ae,sumw[xx]=af,sz[x]=ag,sz[xx]=ah;
	}
}
inline int read() {
	int x=0;
	char c=getchar();
	while(c<'0'||c>'9') c=getchar();
	while(c>='0'&&c<='9') {
		x=(x<<3)+(x<<1)+(c^48);
		c=getchar();
	}
	return x;
}
signed main() { 
	n=read();
	for(int i=1; i<=n; i++) v[i]=read();
	for(int i=1; i<n; i++) {
		int u,v,w;
		u=read(),v=read(),w=read();
		e[u].push_back(node(v,w));
		e[v].push_back(node(u,w));
	}
	calc(1,0);
	dfs(1);
	for(int i=1; i<=n; i++) sum=(sum+anss[i])%p;
	cout<<sum;
	return 0;
} 
2024/3/25 21:32
加载中...