牛客练习赛97_D 月之暗面

原题链接

题目描述

给出一棵 n 个点的树，有 x 种普通颜色，y 种特殊颜色

现在要给树上的每个节点染色，普通颜色染色没有限制，但两个相邻的节点不能染相同颜色的特殊颜色

求染色方案数，答案对 998244353 取模。

输入样例

2 1 2
1 2

输出样例

7

---

算法

(快速幂求逆元 + 树形dp)

本题首先分析题意，有 n 个节点其中每个节点都可以染成普通颜色或是特殊颜色 ，本题是要求染色方案。本质上是搜索所有可能，然而本题 n 的数据范围是1e6 ，且整个染色过为树状结构, 所以可以用树形DP 进行优化，最后只要将根节点的两种状态相加即可。

状态划分： dp[N][2]，以第 i 个节点为根节点，染 j 的颜色的染色方案数（染色情况只有两种，0表示普通颜色，1表示特殊颜色）

状态计算：
分成两种染色情况来写状态转移方程：（利用集合的分部乘法原理与分类加法原理）

$dp[u][0]=dp[u][0]\ast (dp[j][0]+dp[j][1])$
u的普通颜色可以由任何颜色转移

$dp[u][1]=(dp[u][1]\ast (dp[j][0]+(y-1)\ast dp[j][1]\ast qmi(y,mod-2))$

u的特殊颜色只能由普通颜色转移或者由 y个特殊颜色中选择y - 1个进行转移，即为 $dp[u][1]=(dp[u][1]\ast (dp[j][0]+(y-1)/y\ast dp[j][1])$ ,然而，这里的(y - 1) / y 由于取mod 的值为 998244353（为质数）， 所以可以将 $y^{-1}$ 替换成 y 的逆元 qmi(y,mod - 2)

qmi（y,mod - 2）表示用快速幂求逆元
能代替的原因：当 a * x ≡ 1 (mod p),且a 与 p 互质，x即为a的逆元= $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;
}