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} + s z_{v} \times a_{u})$

其中 $a_{u}$ 为题中所给， $b_{u}$ 表示 $a_{u}$ 的位数， $s z_{u}$ 为子树大小。

考虑如何换根，也就是计算 $f a_{v}$ 子树对 $v$ 的贡献。这里将 $f a_{v}$ 写为 $u$ ，有：

$f_{v} = f_{v} + (f_{u} - f_{v} \times b_{u} - s z_{v} \times a_{u}) \times b_{u} + (s z_{1} - s z_{v}) \times a_{u}$

答案 $a n s = \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;
}