P9437 『XYGOI round1』一棵树 题解

赛时一眼换根 dp，然后调调调了大概 1h+。
题目传送门

什么是换根 dp

在大多数树形 dp 中，我们只考虑对根的贡献，而一部分题目需要算出对所有点的贡献，一个比较显然的做法是对每个点都跑一次树形 dp，但是大大增加了时间复杂度，是我们不能接受的。

树形 dp 中的换根 dp 问题又被称为二次扫描，通常不会指定根结点，并且根结点的变化会对一些值，例如子结点深度和、点权和等产生影响。

换根 dp 的常用套路钦定某一个点为根，然后从根开始进行两次 dfs，第一次记录自己子树内中的节点对自己的贡献，第二次记录不属于自己子树中的点对自己的贡献。

Solution

不妨钦定 1 为根。
考虑令 $f_u=\sum _{v\in subtree(u)}w(v,u)$，$g_u=(\sum _{v\notin subtree(u)}w(v,u))+w_u$。为什么要加上 $w_u$？ 是为了方便转移（不加好像也不是很麻烦）。

首先我们要实现一个位数的函数。

LL mylog(LL v) {
	if(v == 0) return 10;
	int res = 0;
	while(v) res ++, v /= 10;
	return pow(10, res);
} 

计算 $f$ 数组

考虑从儿子 $v$ 转移到 $u$，显然为 $f_v\times mylog(w_u)+si{z}_v\times w_u$。

得到 $f_u=(\sum _{v\in subtree(u)}{f}_v)\times mylog(w_u)+w_u\times si{z}_u$ 。

void dfs1(int u, int fa) {
	siz[u] = 1;
	for(auto v : e[u])
		if(v != fa)	{
			dfs1(v, u);
			siz[u] += siz[v];
			f[u] = (f[u] + f[v]) % mod;
		}
	f[u] = (f[u] * mylog(w[u]) % mod + w[u] * siz[u] % mod) % mod;
}

计算 $g$ 数组

如何计算子树外的点对自己的贡献呢？
不妨先想想从根节点转移到第二层节点。

$1$ 转移到 $u$。

先算出除 $subtree(u)$ 外的点对 $1$ 的贡献，显然为 $f_1$ 减去 $u$ 对 $1$ 的贡献 $f_u\ast mylog(w_1)+si{z}_u\ast w_1$。

再考虑 $u$ 对儿子 $v$ 的贡献，为 $f_u+g_u$ 减去 $v$ 对 $u$ 的贡献，再把这些继续用计算 $f$ 的方法转移到 $v$ 上去。

void dfs2(int u, int fa) {
	for(auto v : e[u])
		if(v != fa)	{
			g[v] = ((f[u] + g[u] - val(v, u) - w[u]) % mod * mylog(w[v]) % mod + (siz[1] - siz[v] + 1) % mod * w[v] % mod) % mod;
			dfs2(v, u);
		}
}

AC code

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 1e6 + 5, INF = 0x3f3f3f3f;
const LL mod = 998244353;
int n, m;
LL w[N];
vector<int> e[N];
LL f[N], g[N], siz[N];
LL mylog(LL v) {
	if(v == 0) return 10;
	int res = 0;
	while(v) res ++, v /= 10;
	return pow(10, res);
} 
LL val(int u, int v) {
	return (f[u] * mylog(w[v]) + siz[u] * w[v]) % mod;
}
void dfs1(int u, int fa) {
	siz[u] = 1;
	for(auto v : e[u])
		if(v != fa)	{
			dfs1(v, u);
			siz[u] += siz[v];
			f[u] = (f[u] + f[v]) % mod;
		}
	f[u] = (f[u] * mylog(w[u]) % mod + w[u] * siz[u] % mod) % mod;
}
void dfs2(int u, int fa) {
	for(auto v : e[u])
		if(v != fa)	{
			g[v] = ((f[u] + g[u] - val(v, u) - w[u]) % mod * mylog(w[v]) % mod + (siz[1] - siz[v] + 1) % mod * w[v] % mod) % mod;
			dfs2(v, u);
		}
}
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> n;
    for(int i = 1; i <= n; i ++) cin >> w[i];
    for(int i = 2; i <= n; i ++) {
    	int p;
    	cin >> p;
    	e[p].push_back(i), e[i].push_back(p);
  	}
  	dfs1(1, 0);
  	g[1] = w[1];
  	dfs2(1, 0);
  	LL res = 0;
  	for(int i = 1; i <= n; i ++)
  		res = (res + f[i] + g[i] - w[i]) % mod;
  	cout << res << '\n';
    return 0;
}