[AtCoder Educational DP Contest] V - Subtree（树形dp + 前缀积/后缀积）

problem

luogu

给一棵树，对每一个节点染成黑色或白色。

对于每一个节点，求强制把这个节点染成黑色的情况下，所有的黑色节点组成一个联通块的染色方案数，答案对 $M$ 取模。

$1\le n\le 1e5,2\le M\le 1e9$。

solution

换根 $dp$。

设 $f(i):$ 以 $i$ 为根的子树，$i$ 为黑色的合法方案数。（全局是以 $1$ 为根的）

则有：$f(u)={\prod }_{v\in so{n}_u}(f(v)+1)$。加一是整棵子树全为白色的方案数。

设 $g(i):$ 以 $i$ 为根的子树，$i$ 为黑色的合法方案数。（这里的子树指代的是 $i$ 的祖先及兄弟，即去掉 $i$ 原来子树的剩余部分）

显然答案为 $f(i)\ast g(i)$。

考虑如何求出 $g(i)$。

对于一对父子关系 $f(u)=\prod (f(son)+1)$，其中就有儿子提供的一个贡献 $f(v)+1$，我们需要去掉。

即，$g(v)=g(u)\ast \frac{\prod (f(son)+1)}{f(v)+1}$。

但是 $M$ 并不是质数，所以不能求逆元。

我们可以对每个 $u$ 维护出儿子贡献 $\prod (f(v)+1)$ 的前缀积和后缀积。

只要知道某个儿子在其儿子中的编号即可。

具体可见代码实现。

code

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define maxn 100005
vector < int > pre[maxn], suf[maxn];
vector < int > G[maxn];
int f[maxn], g[maxn];
int n, mod;
void dfs1( int u, int fa ) {
	f[u] = 1;
	for( int v : G[u] )
		if( v == fa ) continue;
		else dfs1( v, u ), f[u] = f[u] * (f[v] + 1) % mod;
	pre[u].resize( G[u].size() + 1, 1 );
	suf[u].resize( G[u].size() + 1, 1 );
	for( int i = 1;i < G[u].size();i ++ ) {
		pre[u][i] = pre[u][i - 1];
		if( G[u][i - 1] ^ fa ) 
			pre[u][i] = pre[u][i] * (f[G[u][i - 1]] + 1) % mod;
	}
	for( int i = G[u].size() - 2;i >= 0;i -- ) {
		suf[u][i] = suf[u][i + 1];
		if( G[u][i + 1] ^ fa )
			suf[u][i] = suf[u][i] * (f[G[u][i + 1]] + 1) % mod;
	}
}
void dfs2( int u, int fa ) {
	for( int i = 0;i < G[u].size();i ++ ) {
		int v = G[u][i];
		if( v == fa ) continue;
		g[v] = pre[u][i] * suf[u][i] % mod * g[u] % mod + 1;//加1是父亲代表的子树全为白 
		dfs2( v, u );
	}
}
signed main() {
	scanf( "%lld %lld", &n, &mod );
	for( int i = 1, u, v;i < n;i ++ ) {
		scanf( "%lld %lld", &u, &v );
		G[u].push_back( v );
		G[v].push_back( u );
	}
	dfs1( 1, 0 );
	g[1] = 1;
	dfs2( 1, 0 );
	for( int i = 1;i <= n;i ++ ) printf( "%lld\n", g[i] * f[i] % mod );
	return 0;
}