Codechef Painting Tree 题解 [ 蓝 ] [ 树形 DP ] [ 概率期望 ] [ 分类讨论 ]

Painting Tree

若干个月前模拟赛切的题，当时写了 3h+，被细节恶心坏了，遂记之。

题意可以转化为求树上存在相交链的期望时间。

考虑如何计算这个期望。显然我们可以枚举选取链的个数，根据期望的定义式来算：

\[E(X) = \sum_{i = 1}^{n}P(Len = i)\times i \]

因为相交链出现在最后一个不交链时刻的下一时刻，所以问题就被我们转化为：对于每个 \(i\)，求出最终不交链的数量 \(= i\) 的概率。

接下来就可以树形 DP 求解树上选择 \(\bm k\) 条链，使其不相交的方案数了。

观察链的性质，我们定义 \(dp_{u, v, 0/1/2}\) 分别表示节点 \(u\)，体积为 \(v\)，\(u\) 上不挂链 / 挂了链不可以延伸 / 挂了链可以延伸这三种状态。并且钦定一条链的体积被计算于它被截断的时刻。

对于转移，我们需要处理以下几种情况：

• 自己和儿子都不挂可延伸的链。
• 自己和某个儿子可延伸的链合成一个新链。
• 从某个儿子处接上可延伸的链。

一共八种转移方程，全部都需要转移，细节很多。

用树形背包的方式进行合并，即可做到时间复杂度 \(O(n^2)\)。

最后统计答案的时候，注意选择链的方案总数为 \((\dfrac{n(n+1)}{2})^i\times i!\)。

#include <bits/stdc++.h>
#define fi first
#define se second
#define eb(x) emplace_back(x)
#define pb(x) push_back(x)
#define lc(x) (tr[x].ls)
#define rc(x) (tr[x].rs)
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef long double ldb;
using pi = pair<int, int>;
const ll mod = 998244353;
const int N = 2005;
int n;
int dp[N][N][3], sz[N], f[N][3], ans; // 0 : 不挂链 ; 1 : 挂链不可延伸 ; 2 : 挂链可延伸
vector<int> g[N];
ll qpow(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 merge(int u, int v)
{
    memcpy(f, dp[u], sizeof(f));
    memset(dp[u], 0, sizeof(dp[u]));
    for(int i = 0; i <= sz[u]; i++)
    {
        for(int j = 0; j <= sz[v]; j++)
        {
            // 自己和儿子都不挂可延伸的链
            dp[u][i + j][0] = (dp[u][i + j][0] + 1ll * f[i][0] * dp[v][j][0] % mod) % mod;
            dp[u][i + j][0] = (dp[u][i + j][0] + 1ll * f[i][0] * dp[v][j][1] % mod) % mod;
            // 自己和某个儿子可延伸的链合成一个新链
            dp[u][i + j + 1][1] = (dp[u][i + j + 1][1] + 1ll * f[i][2] * dp[v][j][2] % mod) % mod;
            dp[u][i + j][1] = (dp[u][i + j][1] + 1ll * f[i][1] * dp[v][j][0] % mod) % mod;
            dp[u][i + j][1] = (dp[u][i + j][1] + 1ll * f[i][1] * dp[v][j][1] % mod) % mod;
            // 从某个儿子处接上可延伸的链
            dp[u][i + j][2] = (dp[u][i + j][2] + 1ll * f[i][2] * dp[v][j][0] % mod) % mod;
            dp[u][i + j][2] = (dp[u][i + j][2] + 1ll * f[i][2] * dp[v][j][1] % mod) % mod;
            dp[u][i + j][2] = (dp[u][i + j][2] + 1ll * f[i][0] * dp[v][j][2] % mod) % mod;
        }
    }
    sz[u] += sz[v];
}
void dfs(int u, int fa)
{
    // 分别对应 不挂链、挂了链但是立刻被截断、挂了可延伸的链 的情况。
    dp[u][0][0] = dp[u][1][1] = dp[u][0][2] = sz[u] = 1;
    for(auto v : g[u])
    {
        if(v == fa) continue;
        dfs(v, u);
        merge(u, v);
    }
}
int main()
{
    // freopen("tree.in", "r", stdin);
    // freopen("tree.out", "w", stdout);
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    cin >> n;
    for(int i = 1; i < n; i++)
    {
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1, 0);
    ll npath = (1ll * (n + 1) * n / 2) % mod, jc = 1;
    for(int i = 1; i <= n + 1; i++)
    {
        if(i - 1) jc = (jc * (i - 1)) % mod;
        ans = (ans + 1ll * (dp[1][i - 1][0] + dp[1][i - 1][1]) % mod * qpow(qpow(npath, i - 1), mod - 2) % mod * jc % mod) % mod;
    }
    cout << ((ans - 1) % mod + mod) % mod;
    return 0;
}