原题链接
题目描述
给出一棵 n 个点的树,有 x 种普通颜色,y 种特殊颜色
现在要给树上的每个节点染色,普通颜色染色没有限制,但两个相邻的节点不能染相同颜色的特殊颜色
求染色方案数,答案对 998244353 取模。
输入样例
2 1 2
1 2
输出样例
7
算法
(快速幂求逆元 + 树形dp)
本题首先分析题意,有 n 个节点其中每个节点都可以染成普通颜色或是特殊颜色 ,本题是要求染色方案。本质上是搜索所有可能,然而本题 n 的数据范围是1e6 ,且整个染色过为树状结构, 所以可以用树形DP 进行优化,最后只要将根节点的两种状态相加即可。
状态划分: dp[N][2],以第 i 个节点为根节点,染 j 的颜色的染色方案数(染色情况只有两种,0表示普通颜色,1表示特殊颜色)
状态计算:
分成两种染色情况来写状态转移方程:(利用集合的分部乘法原理与分类加法原理)
d
p
[
u
]
[
0
]
=
d
p
[
u
]
[
0
]
∗
(
d
p
[
j
]
[
0
]
+
d
p
[
j
]
[
1
]
)
dp[u][0] = dp[u][0] * (dp[j][0] + dp[j][1])
dp[u][0]=dp[u][0]∗(dp[j][0]+dp[j][1])
u的普通颜色可以由任何颜色转移
d p [ u ] [ 1 ] = ( d p [ u ] [ 1 ] ∗ ( d p [ j ] [ 0 ] + ( y − 1 ) ∗ d p [ j ] [ 1 ] ∗ q m i ( y , m o d − 2 ) ) dp[u][1] = (dp[u][1] * (dp[j][0] + (y - 1) * dp[j][1] * qmi(y,mod - 2)) dp[u][1]=(dp[u][1]∗(dp[j][0]+(y−1)∗dp[j][1]∗qmi(y,mod−2))
u的特殊颜色只能由普通颜色转移或者由 y个特殊颜色中选择y - 1个进行转移,即为 d p [ u ] [ 1 ] = ( d p [ u ] [ 1 ] ∗ ( d p [ j ] [ 0 ] + ( y − 1 ) / y ∗ d p [ j ] [ 1 ] ) dp[u][1] = (dp[u][1] * (dp[j][0] + (y - 1) / y * dp[j][1] ) dp[u][1]=(dp[u][1]∗(dp[j][0]+(y−1)/y∗dp[j][1]) ,然而,这里的(y - 1) / y 由于取mod 的值为 998244353(为质数), 所以可以将 y − 1 y^{-1} y−1 替换成 y 的逆元 qmi(y,mod - 2)
qmi(y,mod - 2)表示用快速幂求逆元
能代替的原因:当 a * x ≡ 1 (mod p),且a 与 p 互质,x即为a的逆元= a − 1 a^{-1} a−1 ,通过费马小定理可以得到逆元也同样可以用快速幂求
对逆元具体不 在赘述,可以参考
C++ 代码
#include <iostream>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;
typedef long long ll;
const ll N = 1e6 + 10,mod = 998244353;
ll dp[N][2];//以第i个节点为根节点,染j的颜色的染色方案数,树形dp
ll n,x,y,res;
vector<ll> g[N];
bool has_rooted[N];
ll qmi(ll a,ll b)
{
ll res = 1;
while (b) {
if (b & 1) res = res * a % mod;
b >>= 1;
a = a * a % mod;
}
return res;
}
void dfs(int u) //枚举每个节点的状态,并利用状态转移方程计算出每个状态所有可能的情况
{
dp[u][0] = x,dp[u][1] = y;//基本情况
for (int i = 0;i < g[u].size();i ++ ) {
int j = g[u][i];
dfs(j);
dp[u][0] = (dp[u][0] * ((dp[j][0] + dp[j][1]) % mod)) % mod;
dp[u][1] = (dp[u][1] * (dp[j][0] + (y - 1) * dp[j][1] % mod * qmi(y,mod - 2) % mod)) % mod;
}
}
int main()
{
scanf("%d%d%d",&n,&x,&y);
for (int i = 1;i <= n - 1;i ++ ) {
int a,b;
scanf("%d%d",&a,&b);
g[a].push_back(b);
has_rooted[b] = true;
}
int root = 1;
while (has_rooted[root]) root ++;
dfs(root);
cout << (dp[root][0] + dp[root][1]) % mod << '\n';
return 0;
}