[生成函数] Codeforces 891E Lust

题目大意

给你一个长为 \(n\) 的序列 \(\{a_n\}\)，和一个初值为0的变量 \(s\)，要求你重复以下操作 \(k\) 次：

• 在 \(1,2,\cdots,n\) 中等概率随机选择一个 \(x\)
• 令 \(s\) 加上 \(\prod_{i\neq x}a_i\)
• 令 \(a_i\) 减一

求 \(k\) 次操作后 \(s\) 的期望。

\(1\leq n\leq 5000,1\leq k\leq 10^9,0\leq a_i\leq 10^9\)。

题解

假设当前令 \(a_x\) 减一，那么

\[ s+=\prod_{i\neq x}a_i=(a_x-(a_x-1))\prod_{i\neq x}a_i\\ =\prod_ia_i-(a_x-1)\prod_{i\neq x} a_i \]

发现 \(s\) 加上的是数列操作前后整个数列乘积的差值。显然操作多次的贡献就是还没有操作时数列的乘积减去操作结束后数列的乘积。设 \(a_i\) 被操作了 \(b_i\) 次，变成了 \(a_i-b_i\)，则有

\[s=\prod_{i=1}^na_i-\prod_{i=1}^n(a_i-b_i) \]

因此我们只需要求得 \(\prod_{i=1}^n(a_i-b_i)\) 的期望。

对 \(\{a_n\}\) 操作 \(k\) 次有 \(n^k\) 次种方案，对于每一种 \(\{b_n\}\)，又有 \(\frac{k!}{b_1!b_2!\cdots b_n!}\) 种操作方案，所以有

\[E\left(\prod_{i=1}^n(a_i-b_i)\right)=\frac{1}{n^k}\sum_{\{b_n\}}\frac{k!}{\prod_{i=1}^nb_i!}\prod_{i=1}^n(a_i-b_i)\\ =\frac{k!}{n^k}\sum_{\{b_n\}}\prod_{i=1}^n\frac{a_i-b_i}{b_i!} \]

设 \(a_i\) 被操作 \(j\) 次的贡献是 \(\frac{a_i-j}{j!}\)，设这个生成函数为

\[f_i(x)=\sum_{j=0}^\infty\frac{a_i-j}{j!} x^j \]

则

\[\sum_{\{b_n\}}\prod_{i=1}^n\frac{a_i-b_i}{b_i!}=[x^k]\prod_{i=1}^n f_i(x)=[x^k]\prod_{i=1}^n\sum_{j=0}^\infty\frac{a_i-j}{j!}x^j\\ =[x^k]\prod_{i=1}^n\left(a_i\sum_{j=0}^\infty\frac{x^j}{j!}-x\sum_{j=1}^\infty\frac{x^{j-1}}{(j-1)!}\right)\\ =[x^k]\prod_{i=1}^n\left(a_i\sum_{j=0}^\infty\frac{x^j}{j!}-x\sum_{j=0}^\infty\frac{x^j}{j!}\right)\\ =[x^k]\prod_{i=1}^n(a_i-x)e^x\\ =[x^k]e^{nx}\prod_{i=1}^n(a_i-x) \]

因为 \(n\leq 5000\)，所以对于 \(\prod_{i=1}^n(a_i-x)\)，可以 \(O(n^2)\) 暴力求出。设 \(g(x)=\prod_{i=1}^n(a_i-x)\)，则我们要求 \([x^k]e^{nx}g(x)\)。不妨再把 \(e^{nx}\) 展开，有

\[[x^k](e^{nx}g(x))=[x^k]\left(g(x)\sum_{i=0}^\infty\frac{n^i}{i!}x^i\right)\\ =\sum_{i=0}^n\left([x^i]g(x)\times\frac{n^{k-i}}{(k-i)!}\right) \]

回代得

\[E\left(\prod_{i=1}^n(a_i-b_i)\right)=\frac{k!}{n^k}\sum_{i=0}^n\left([x^i]g(x)\times\frac{n^{k-i}}{(k-i)!}\right)\\ =\sum_{i=0}^n\left([x^i]g(x)\times\frac{k^{\underline{i}}}{n^i}\right) \]

该式可以 \(O(n)\) 求出。最后，

\[E(s)=\prod_{i=1}^na_i-\sum_{i=0}^n\left([x^i]g(x)\times\frac{k^{\underline{i}}}{n^i}\right) \]

因此本题的时间复杂度为 \(O(n^2)\)。

Code

#include <bits/stdc++.h>
using namespace std;

#define RG register int
#define LL long long

const LL MOD = 1000000007LL;
LL a[5005], g[5005];
int n, k;

LL qpow(LL b, LL n) {
    LL x = 1, Power = b % MOD;
    while (n) {
        if (n & 1) x = x * Power % MOD;
        Power = Power * Power % MOD;
        n >>= 1;
    }
    return x;
}

LL calc() {
    g[0] = 1;
    for (int i = 1;i <= n;++i) {
        for (int j = n;j >= 1;--j)
            g[j] = ((g[j] * a[i] % MOD - g[j - 1]) % MOD + MOD) % MOD;
        g[0] = g[0] * a[i] % MOD;
    }
    LL ninv = qpow(n, MOD - 2), temp = 1, ans = 0;
    for (int i = 0;i <= n;++i) {
        ans = (ans + g[i] * temp % MOD) % MOD;
        temp = temp * ninv % MOD * (k - i) % MOD;
    }
    temp = 1;
    for (int i = 1;i <= n;++i) temp = temp * a[i] % MOD;
    ans = ((temp - ans) % MOD + MOD) % MOD;
    return ans;
}

int main() {
    scanf("%d%d", &n, &k);
    for (int i = 1;i <= n;++i)
        scanf("%I64d", &a[i]);
    printf("%I64d\n", calc());

    return 0;
}