AT4543 Subtrees

题目链接

洛谷：AT4543

AT：AT4543/DP_V

---

题意简述

给定一棵树，用黑白两色染色。

对于每一个点，求强制给该点染上黑色时，整棵树上的黑点构成一个连通块的染色方案数。

答案取模，模数输入给出。

---

题解

题目要求我们求出对每一个点强制染黑的情形的答案，故考虑采用换根 DP。

先计算钦定 $1$ 号点为根并染黑的方案数。

设 $f_u$ 表示将 $u$ 号点染黑，且其子树内黑点构成连通块的方案数。易见，对于其每一个子节点 $v$，都有染黑和不染黑两种选择：染黑则方案数为 $f_v$；不染黑则 $v$ 的整棵子树都只能为白，方案数为 $1$。故

$$f_u=\prod _{v\in {\mathrm{s}\mathrm{o}\mathrm{n}}_u}(f_v+1)$$

当 $u$ 为叶子时，显然 $f_u=1$，与初始化要求相同。

现在考虑换根。不难理解，换根时消去 $u$ 对 ${\mathrm{f}\mathrm{a}}_u$ 的贡献，再将 ${\mathrm{f}\mathrm{a}}_u$ 的贡献乘到 $u$ 上。

• 消去 $u$ 对 ${\mathrm{f}\mathrm{a}}_u$ 的贡献：

  $$f_{{\mathrm{f}\mathrm{a}}_u}^{\prime }=\frac{f_{{\mathrm{f}\mathrm{a}}_u}}{f_u+1}$$

• 将 ${\mathrm{f}\mathrm{a}}_u$ 的贡献乘到 $u$ 上：

  $$f_u^{\prime }=f_u\times (f_{{\mathrm{f}\mathrm{a}}_u}^{\prime }+1)$$

不过出现了一个小问题：模数不保证是质数，所以不能用直接乘逆元的方式来消去贡献。

这里，对于每一个点 $u$，处理出 $(f_{v\in {\mathrm{s}\mathrm{o}\mathrm{n}}_u}+1)$ 的前缀积和后缀积，即可解决求出消去一个子树贡献后的答案的问题。

---

代码

#include<bits/stdc++.h>
#define eb emplace_back

using namespace std;

namespace acah
{
	using ll = long long;
	
	constexpr int maxn = 1e5 + 7;
	
	int N, p;
	vector<int> T[maxn], sn[maxn];
	ll f[maxn], ans[maxn];
	vector<ll> pf[maxn], sf[maxn];
	
	void dp(int u, int fa)
	{
		f[u] = 1;
		
		for(int v : T[u]) {
			if(v == fa) continue;
			sn[u].eb(v);
			dp(v, u);
			(f[u] *= (f[v] + 1)) %= p;
		}
		
		pf[u].eb(1); // pf[u][0] = 1
		for(int i = 0, bd = sn[u].size(); i < bd; i++)
			pf[u].eb(pf[u][i] * (f[sn[u][i]] + 1) % p);
		
		sf[u].eb(1);
		for(int i = sn[u].size() - 1, c = 0; i >= 0; i--, c++)
			sf[u].eb(sf[u][c] * (f[sn[u][i]] + 1) % p);
		reverse(sf[u].begin(), sf[u].end());
	}
	
	void solve(int u, ll fv)
	{
		ans[u] = f[u] * (fv + 1) % p;
		
		for(int i = 0, bd = sn[u].size(); i < bd; i++) {
			int v = sn[u][i];
			ll l = pf[u][i], r = sf[u][i + 1];
			solve(v, l * r % p * (fv + 1) % p);
		}
	}
	
	int work()
	{
		scanf("%d%d", &N, &p);
		for(int i = 1, u, v; i < N; i++) {
			scanf("%d%d", &u, &v);
			T[u].eb(v), T[v].eb(u);
		}
		
		dp(1, 0), solve(1, 0);
		
		for(int i = 1; i <= N; i++)
			printf("%lld\n", ans[i]);
		
		return 0;
	}
}

int main() {return acah::work();}

这里，因为要使用儿子的编号，所以用基于范围的 for 循环并不如枚举下标简便。