P9437 『XYGOI round1』一棵樹

Fire_Raku發表於2024-04-16
Go

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;
}

相關文章