题解：2024牛客暑期多校第5场——C [骰]

英文题干

wzj is working on a cloze test. For those who are not familiar with cloze tests, here is an introduction:

A cloze test consists of n questions, each with exactly m options. A sequence of n positive integers between 1 and m is called an answer set.

The standard answer to the cloze test is a specific answer set. The score of an answer set is defined as the number of positions where it matches the standard answer. An answer set is called balanced if each integer from 1 to m appears exactly the same number of times (in this problem, n is guaranteed to be a multiple of m). The standard answer is always balanced.

Here is some information about wzj:

wzj has a clear understanding of his own level in cloze tests. He knows that he has a probability $p_i$ of scoring exactly i. This means he has a probability $p_0$ of getting all wrong, a probability $p_1$ of getting exactly 1 correct, ..., and a probability $p_n$ of getting all correct. wzj selects among answers with the same score uniformly at random. That is, wzj's process of completing the cloze test is as follows: he first randomly determines his score according to the probabilities $p_i$ , and then randomly selects one answer set among those that achieve this score.

Given this background:

One day, wzj completed a cloze test and was pleasantly surprised to find that his answer was balanced. Calculate the expected value of his score. Output the result modulo 998244353.

Input:

The first line contains two positive integers n and m, indicating the number of questions and the number of options for each question. It is guaranteed that n is a multiple of m.

The second line contains $n 1$ non-negative integers $p_0,p_1,...,p_n$ , representing the probability of wzj scoring each possible score. The sum of $p_i$ is guaranteed to be 1 modulo 998244353.

Output:

A single integer representing the expected value of wzj's score modulo 998244353.

中文题干

wzj正在进行一个填空测试。对于那些不熟悉填空测试的人，这里有一个介绍：

填空测试由n个问题组成，每个问题有m个选项。n个介于1和m之间的正整数序列称为答案集。(俗称完形填空)

填空测试的标准答案是一个特定的答案集。答案集的得分定义为与标准答案匹配的位置数量。如果每个从1到m的整数出现的次数完全相同，则称答案集为平衡的（在此问题中，n保证是m的倍数）。标准答案总是平衡的。

以下是有关wzj的一些信息：

wzj清楚自己在填空测试中的水平。他知道他有概率 $P_i$ 能够得到正好i分。这意味着他有概率 $P_0$ 得分全错，概率 $P_1$ 得分正好1分，...，以及概率 $P_n$ 得分全对。wzj在相同得分的答案中均匀随机选择答案。也就是说，wzj完成填空测试的过程如下：他首先根据概率 $P_i$ 随机确定自己的得分，然后在那些实现这个得分的答案集中随机选择一个答案集。

鉴于这些背景：

一天，wzj完成了一个填空测试，并惊喜地发现他的答案是平衡的。计算他的得分的期望值。输出结果需对998244353取模。

思路

(Inspired by 命题组)

1.先做概率的数学模型分析：

考虑所有配平的答案里，做对k个题的方案数为 $f_k$ 。
所有做对的k个题中，配平的情况为 $g_k$ ，则 $g_k$ = $\frac{f_k}{C^{k}_{n}\times (m-1)^{n-k}}$
利用 $f_k$ 和 $P_k$ 可由Bayes写出答案： $ans$ = $\frac{\sum{g_iP_ii}}{\sum{g_iP_i}}$

2.下面求 $f_k$ :

不妨设答案为 aaabbbcccddd...
计算 $f_k$ 考虑容斥，先计算钦定 $k$ 个位置是正确的方案数，再进行一次二项式反演
$dp_{i,j}$ 表示考虑了前 $i$ 个选项（即前 $\frac{in}{m}$ 个题），钦定了 $j$ 个题是正确的。
转移枚举第 $i$ 个选项有 $c_i$ 个题是正确的，其余 $n/m-c_i$ 个题要扔到最后随便选。
最后随便选的组合数是 $\frac{(n-\sum{c_i})!}{({n}/{m}-c_1)!\cdot({n}/{m}-c_2)!\dots({n}/{m}-c_m)!}$ ，可以把 $\frac{1}{(n/m-c_i)!}$ 的系数扔到dp中计算。

AC代码

时间复杂度：O( $n^2$ ) (状态数 $O(nm)$ ，转移 $O(n/m)$ )

// 视频题解：https://www.bilibili.com/video/BV1i2421Z79K/?spm_id_from=333.999.0.0&vd_source=2
// 文字题解：https://blog.nowcoder.net/n/dd3b4d893e904ed1913e21ba2f446eaf
// 概率论 + dp + 容斥

#include <bits/stdc++.h>
#define int long long
using namespace std;
const int maxn = 2010;
const int mod = 998244353;

int n, m;
int p[maxn];
int f[maxn], g[maxn];
int dp[maxn][maxn];

// 【求组合数和排列数】
int C[maxn][maxn];
int A[maxn][maxn];
void getCombination()
{
    C[0][0] = 1;
    for (int i = 0; i <= n; i++)
    {
        C[i][0] = 1;
        C[i][i] = 1;
        A[i][0] = 1;
        for (int j = 1; j < i; j++)
        {
            C[i][j] = (C[i - 1][j] + C[i - 1][j - 1]) % mod;
        }
        for (int j = 1; j <= i; j++)
        {
            A[i][j] = A[i][j - 1] * (i - j + 1) % mod;
        }
    }
}
// 【带余快速幂】
int qpow(int a, int b)
{
    int res = 1;
    while (b)
    {
        if (b & 1)
            res = res * a % mod;
        a = a * a % mod;
        b >>= 1;
    }
    return res % mod;
}

// 【求逆元】
int inverse(int x)
{
    return qpow(x, mod - 2);
}

signed main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);

    cin >> n >> m;
    for (int i = 0; i <= n; i++)
        cin >> p[i];

    getCombination();

    // 计算dp[m][k]: 一共m题，至少做对k个题的方案数
    dp[0][0] = 1;
    for (int i = 1; i <= m; i++)
    {
        for (int j = 0; j <= n; j++)
        {
            // 如果前i-1个选项中还没钦定至少j个正确，则跳过
            if (!dp[i - 1][j])
                continue;
            // 枚举第i个选项还有k个是正确的
            for (int k = 0; k <= n / m; k++)
            {
                // dp[i][j+k] += dp[i-1][j] * C[n/m][k] / (n/m-k)!
                dp[i][j + k] = (dp[i][j + k] + dp[i - 1][j] * C[n / m][k] % mod * inverse(A[n / m - k][n / m - k]) % mod) % mod;
            }
        }
    }

    // 把"随便选"的分子补上，并进行容斥运算得到f[k]
    for (int i = n; i >= 0; i--)
    {
        dp[m][i] = dp[m][i] * A[n - i][n - i] % mod;
        // 容斥减去重复的
        for (int j = i + 1; j <= n; j++)
            dp[m][i] = (dp[m][i] - dp[m][j] * C[j][i] % mod + mod) % mod;
    }

    // 计算g[k]
    for (int i = 0; i <= n; i++)
        g[i] = dp[m][i] % mod * inverse(C[n][i] * qpow(m - 1, n - i) % mod) % mod;

    // 计算ans
    int fz = 0, fm = 0;
    for (int i = 0; i <= n; i++)
    {
        fz = (fz + p[i] * g[i] % mod * i) % mod;
        fm = (fm + p[i] * g[i]) % mod;
    }
    cout << fz * inverse(fm) % mod;

    return 0;
}
// Inspired by 命题组