题意
给定长度为 \(n\) 的数组 \(a\) 和一个整数 \(k\) ，执行下面的代码：
long long ans = 0; //定义一个初始值为0的长整型变量
for(int i = 1; i <= k; i++) {
int idx = rnd.next(0, n - 1); //生成一个介于0到n-1的随机数（含0和n-1）
//每个数被选中的概率是相同的
ans += a[idx];
a[idx] -= (a[idx] % i);
}
求代码运行结束后 \(ans\) 的期望值 \(E\) 与 \(n^k\) 相乘的结果，即 \(E \times n^k\)。
（\(1 \le n \le 10^7, 1 \le k \le 17\)）。
题解
考虑计算 \(ans\) 的期望值 \(E\)，最后使之与 \(n^k\) 相乘。设 \(P = \operatorname{lcm}\left(1 \sim k - 1\right)\)，可以看出数组 \(a\) 中的元素 \(a_i\) 最多会被减去 \(a_i \bmod P\)。因此若将 \(a_i\) 拆分为两部分：\(a_i = \left\lfloor\dfrac{a_i}{P}\right\rfloor \times P + a_i \bmod P\)，前半部分的贡献是不变的，为
\[k \times \dfrac{P \times \sum\limits_{i = 0}^{n - 1}\left\lfloor\dfrac{a_i}{P}\right\rfloor}{n}
\]
下面着重考虑后半部分的贡献，因为值域较小，故考虑将其作为 \(\texttt{DP}\) 的状态。设 \(f_{i, j}\) 为在第 \(i\) 次操作后模 \(P\) 为 \(j\) 的期望元素个数。有转移
\[\begin{aligned}
f_{i, j - \left(j \bmod i\right)} &\leftarrow \dfrac{1}{n} \times f_{i - 1, j}\\
f_{i, j} &\leftarrow \left(1 - \dfrac{1}{n}\right) \times f_{i - 1, j}
\end{aligned}\]
计算完成 \(f\) 数组后即可快速计算前半部分的贡献，为
\[\sum\limits_{i = 0}^{k - 1} \dfrac{\sum\limits_j j \times f_{i, j}}{n}
\]
总复杂度 \(\mathcal{O}(n + kP)\)，可以通过本题。
Code
#include <bits/stdc++.h>
typedef long long valueType;
typedef std::vector<valueType> ValueVector;
typedef std::vector<ValueVector> ValueMatrix;
constexpr valueType MOD = 998244353;
template<typename T1, typename T2, typename T3 = valueType>
void Inc(T1 &a, T2 b, const T3 &mod = MOD) {
a = a + b;
if (a >= mod)
a -= mod;
}
template<typename T1, typename T2, typename T3 = valueType>
T1 sub(T1 a, T2 b, const T3 &mod = MOD) {
return a - b < 0 ? a - b + mod : a - b;
}
template<typename T1, typename T2, typename T3 = valueType>
T1 mul(T1 a, T2 b, const T3 &mod = MOD) {
return (long long) a * b % mod;
}
template<typename T1, typename T2, typename T3 = valueType>
void Mul(T1 &a, T2 b, const T3 &mod = MOD) {
a = (long long) a * b % mod;
}
template<typename T1, typename T2, typename T3 = valueType>
T1 pow(T1 a, T2 b, const T3 &mod = MOD) {
T1 result = 1;
while (b > 0) {
if (b & 1)
Mul(result, a, mod);
Mul(a, a, mod);
b = b >> 1;
}
return result;
}
valueType lcm(valueType a, valueType b) {
return a / std::__gcd(a, b) * b;
}
int main() {
valueType N, A0, X, Y, K, M;
std::cin >> N >> A0 >> X >> Y >> K >> M;
ValueVector source(N);
source[0] = A0;
for (valueType i = 1; i < N; ++i)
source[i] = (source[i - 1] * X + Y) % M;
valueType P = 1;
for (valueType i = 2; i < K; ++i)
P = lcm(P, i);
ValueMatrix F(K, ValueVector(P, 0));
for (auto const &iter: source)
++F[0][iter % P];
valueType const InvN = pow(N, MOD - 2), reverseN = sub(1, InvN);
for (valueType k = 1; k < K; ++k) {
for (valueType i = 0; i < P; ++i) {
Inc(F[k][i - i % k], mul(F[k - 1][i], InvN));
Inc(F[k][i], mul(F[k - 1][i], reverseN));
}
}
valueType ans = 0;
for (valueType k = 0; k < K; ++k)
for (valueType i = 0; i < P; ++i)
Inc(ans, mul(F[k][i], i));
for (auto const &iter: source)
Inc(ans, mul(iter / P * P, K));
Mul(ans, InvN);
Mul(ans, pow(N, K));
std::cout << ans << std::endl;
return 0;
}