换根 DP WA5pts 求调
查看原帖
换根 DP WA5pts 求调
386892
REMAC楼主2023/11/16 17:25

rt,怀疑是取模的问题。

rv[i]是 10i10^i。

#include<bits/stdc++.h>
using namespace std;
//dp: 子树内点到u的值和

const int maxn=1e6+10;

vector<int> g[maxn];

#define int long long

const int mod=998244353;

int w[maxn];

int a[maxn],dp[maxn],sz[maxn];

inline int ws(int x) {
	if(x>=1000000) return 7;
	if(x>=100000) return 6;
	if(x>=10000) return 5;
	if(x>=1000) return 4;
	if(x>=100) return 3;
	if(x>=10) return 2;
	return 1;
}

__int128 rv[]= {0,10,100,1000,10000,100000,1000000,10000000};

int qpow(int a,int p) {
	int ret=1;
	while(p>0) {
		if(p&1) ret=(ret*a)%mod;
		a=(a*a)%mod;
		p>>=1;
	}
	return ret;
}

void dfs(int u,int f) {
	sz[u]=1;
	int ww=ws(w[u]);
	dp[u]=(dp[u]+w[u])%mod;
	for(int v:g[u]) {
		if(v==f) continue;
		dfs(v,u);
		dp[u]=(dp[u]+dp[v]*(rv[ww]%mod)%mod+w[u]*sz[v])%mod;
		sz[u]+=sz[v];
	}
}

void dfs2(int u,int f) {
	int ww=ws(w[u]);
	if(u!=1) {
		a[u]=(((a[f]-dp[u]*rv[ws(w[f])]-sz[u]*w[f]%mod)*rv[ww]%mod)%mod+w[u]*(sz[1]-sz[u]+1))%mod;
		for(int v:g[u]) {
			if(v==f) continue;
			a[u]=(a[u]+dp[v]*rv[ww]%mod+w[u]*sz[v]%mod)%mod;
		}
	}
	for(int v:g[u]) {
		if(v==f) continue;
		dfs2(v,u);
	}
}

main() {
	int n;
	cin>>n;
	for(int i=1; i<=n; i++) cin>>w[i];
	for(int i=2; i<=n; i++) {
		int v;
		cin>>v;
		g[i].push_back(v);
		g[v].push_back(i);
	}

	dfs(1,1);
	a[1]=dp[1]%mod;
	dfs2(1,1);
	int ans=0;
	for(int i=1; i<=n; i++) ans=(ans+a[i])%mod;
	cout<<(ans+mod)%mod;
}
2023/11/16 17:25
加载中...