Codeforces 891E Lust 生成函数

题意

现在给你一有n个整数的序列a[]，有一个初始为0的值res，重复下面的过程k次：
“随机选择一个[1,n]之间的下标x，res加上所有满足i≠x的a[i]的乘积，然后将a[x]减去1”
问最后res的期望值，对10^9+7取模
n<=5000
k<=10^9

分析

我们把对a[x]的一次操作的贡献看成是

$$\prod a[i]-\prod a'[i]$$
其中 $a'[i]$表示将 $a[x]$减去1后的数组。
不难发现通过连续的操作后，剩下的项就只剩下 $$\prod a[i]-\prod (a[i]-b[i])$$
其中 $b[i]$表示 $a[i]$被操作的次数。那么我们要求的就是上述式子的期望。
左边的式子显然是常数，那么只考虑右边的式子。
对于b是数组的每一种取值，对答案的贡献显然是 $$\frac{k!\prod(a[i]-b[i])}{n^k\prod (b[i])!}$$
设生成函数 $f_i(x)$表示 $b[i]$取每一种值时对答案的贡献，显然有 $$f_i(x)=\sum_{j>=0}\frac{(a[i]-j)x^j}{j!}$$
化简后得 $$f_i(x)=(a[i]-x)e^x$$
那么我们要求的答案就是多项式 $e^{nx}\prod(a[i]-x)$的第k项系数。
右边可以暴力卷积，显然次数为n，左边展开后可以得到 $$e^{nx}=\sum_{i>=0}\frac{(nx)^i}{i!}$$
然后继续暴力卷积即可。
时间复杂度 $O(n^2)$

代码

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>
using namespace std;

typedef long long LL;

const int N=5005;
const int MOD=1000000007;

int n,f[N],k;

int ksm(int x,int y)
{
    int ans=1;
    while (y)
    {
        if (y&1) ans=(LL)ans*x%MOD;
        x=(LL)x*x%MOD;y>>=1;
    }
    return ans;
}

int dwm(int x,int y)
{
    int ans=1;
    while (y--)
    {
        ans=(LL)ans*x%MOD;
        x--;
    }
    return ans;
}

int main()
{
    scanf("%d%d",&n,&k);
    f[0]=1;int s=1;
    for (int i=1;i<=n;i++)
    {
        int a;scanf("%d",&a);
        s=(LL)s*a%MOD;
        for (int j=i;j>=1;j--) f[j]=((LL)f[j]*a%MOD+MOD-f[j-1])%MOD;
        f[0]=(LL)f[0]*a%MOD;
    }
    int ans=0;
    for (int i=0;i<=min(n,k);i++) (ans+=(LL)ksm(n,k-i)*dwm(k,i)%MOD*f[i]%MOD)%=MOD;
    ans=(LL)ans*ksm(n,(LL)k*(MOD-2)%(MOD-1))%MOD;
    ans=(s+MOD-ans)%MOD;
    printf("%d",ans);
    return 0;
}