P9437 『XYGOI round1』一棵树 （换根dp）

P9437 『XYGOI round1』一棵树

trick+换根dp

对于此类 「将数字顺次写下」 计算贡献的题目，通常按位考虑，并且考虑每个数作为 开头/结尾 时的贡献，方便计算。

因此，我们在这题中考虑每个数作为结尾时的贡献。那么这题就转化成：计算以 \(u\) 为根并且以 \(a_u\) 为结尾的贡献。明显的换根 dp。

首先考虑处理出各个子树中的贡献，设 \(f_u\) 表示在 \(u\) 子树中，以 \(a_u\) 为结尾的贡献。显然有：

\(f_u=a_u+\sum (f_v\times b_u+sz_v\times a_u)\)

其中 \(a_u\) 为题中所给，\(b_u\) 表示 \(a_u\) 的位数，\(sz_u\) 为子树大小。

考虑如何换根，也就是计算 \(fa_v\) 子树对 \(v\) 的贡献。这里将 \(fa_v\) 写为 \(u\)，有：

\(f_v=f_v+(f_u-f_v\times b_u-sz_v\times a_u)\times b_u+(sz_1-sz_v)\times a_u\)

答案 \(ans=\sum f_u\)，复杂度 \(O(n)\)。

#include <bits/stdc++.h>
#define pii std::pair<int, int>
#define fi first
#define se second
#define pb push_back

typedef long long i64;
const i64 iinf = 0x3f3f3f3f, linf = 0x3f3f3f3f3f3f3f3f;
const int N = 1e6 + 10, mod = 998244353;
int n;
i64 ans, c[N], f[N], sz[N], b[N], a[N];
std::vector<int> e[N];
void dfs1(int u, int fa) {
	sz[u] = 1, f[u] = a[u];
	for(auto v : e[u]) {
		if(v == fa) continue;
		dfs1(v, u);
		f[u] = (f[u] + f[v] * b[u] % mod + sz[v] * a[u] % mod) % mod;
		sz[u] += sz[v];
	}
}
void dfs2(int u, int fa) {
	ans = (ans + f[u]) % mod;
	for(auto v : e[u]) {
		if(v == fa) continue;
		f[v] = (f[v] + (f[u] - f[v] * b[u] % mod + mod - sz[v] * a[u] % mod + mod) * b[v] % mod + (sz[1] - sz[v]) * a[v] % mod) % mod;
		dfs2(v, u);
	}
}
void Solve() {
	std::cin >> n;
	c[1] = 10;
	for(int i = 2; i <= 9; i++) c[i] = c[i - 1] * 10;
	for(int i = 1; i <= n; i++) {
		std::cin >> a[i];
		if(!a[i]) b[i] = 10;
		else {
			int cnt = 0;
			i64 now = a[i];
			while(now) {
				now /= 10;
				cnt++;
			}
			b[i] = c[cnt];
		}
	}
	for(int i = 2; i <= n; i++) {
		int x;
		std::cin >> x;
		e[x].pb(i), e[i].pb(x);
	}
	dfs1(1, 0);
	dfs2(1, 0);
	std::cout << ans << "\n";
}
int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    
	Solve();

	return 0;
}