[树形dp] CF1146F Leaf Partition

posted on 2024-06-05 05:25:42 | under | source

尝试找性质，无果。

尝试直接树形 \(\rm dp\)，那么应当是枚举 \(u\) 是否作为 \(\rm lca\) 被选取。

计算不钦定 \(u\) 的方案是容易的，那么钦定呢？我们需要枚举哪些子树 \(v\) 向上延伸，剩下的子树随意取即可。

注意一下，假如只考虑 \(v\) 子树，记 \(v\) 被选取的方案为 \(cnt\)。那么对其父亲 \(u\) 而言，这个 \(cnt\) 其实有两种含义：\(v\) 向上延伸与 \(u\) 相连的方案数、\(v\) 不向上延伸自己做主的方案数。

（其实是在打代码时发现的问题。）

于是树形 \(\rm dp\) 的状态，应该要根据其是否与父亲相连而决定。不难得出定义 \(f_{u,0/1}\) 表示 \(u\) 子树，\(u\) 是否向上延伸（\(u\) 父亲是否被选取）的方案数。

然后对 \(u\) 处理出 \(g_{u,0/1/2}\)，表示有 \(0/1/\ge2\) 个子树向上延伸相连的方案数。

那么转移就是：\(f_{u,0}=g_{u,0}+g_{u,2}\)，\(f_{u,1}=g_{u,1}+g_{u,2}\)。

复杂度 \(O(n)\)。

代码

#include<bits/stdc++.h>
using namespace std;

#define int long long
#define ADD(a, b) a = (a + (b)) % mod
const int N = 2e5 + 5, mod = 998244353;
int n, fa, f[N][2];
vector<int> to[N];

inline void dfs(int u, int fa){
	if(to[u].empty()) {f[u][0] = f[u][1] = 1; return ;}
	int g[5], _g[5]; memset(g, 0, sizeof g), g[0] = 1;
	for(auto v : to[u])
		if(v ^ fa){
			dfs(v, u);
			memcpy(_g, g, sizeof _g), memset(g, 0, sizeof g);
			for(int j = 2; ~j; --j) 
				ADD(g[min(2ll, j + 1)], _g[j] * f[v][1] % mod), ADD(g[j], _g[j] * f[v][0] % mod);
		}
	f[u][0] = (g[0] + g[2]) % mod;
	f[u][1] = (g[1] + g[2]) % mod;
}
signed main(){
	cin >> n;
	for(int i = 2; i <= n; ++i) scanf("%lld", &fa), to[fa].push_back(i);
	dfs(1, 0);
	cout << f[1][0];
	return 0;
}