Atcoder 231

\(G.Balls in Boxes\)(期望，数学，多项式)

题意

有\(N\)个盒子，每个盒子一开始有\(A_i\)个球，可以进行\(K\)次操作，每次操作可以随机选一个盒子加一个球，\(K\)次操作后每个盒子有\(B_i\)个球，令\(S=\prod_{i=1}^nB_i\)，求\(E(S)\)，答案对\(998244353\)取模

\(1\le N\le 1000\) \(1\le K\le 10^9\) \(1\le A_i\le10^9\)

Sol

设\(X_i\)为第\(i\)个盒子被选中的随机变量。则\(E(S)=E(\prod_{i=1}^n(A_i+X_i))\)

取\(n=2\)为例子，\(E(S)=E(A_1A_2)+E(A_1X_2)+E(A_2X_1)+E(X_1X_2)\)

由于\(A_1,A_2\)是常数，所以\(E(S)=A_1A_2+A_1E(X_2)+A_2E(X_1)+E(X_1X_2)\)

\(n=3\)时，\(E(S)=E(X_1X_2X_3)+A_1E(X_2X_3)+A_2E(X_1X_3)+A_3E(X_1X_2)+A_1A_2E(X_3)+A_1A_3E(X_2)+A_2A_3E(X_1)+A_1A_2A_3\)

观察前面的系数为\(1+A_1 + A_2+ A_3+ A_1A_2 + A_1A_3 + A_2A_3+ A_1A_2A_3=(1+A_1)(1+A_2)(1+A_3)\)

记\(S=\prod_{i=1}^n(1+A_ix)\)，\(S(i)\)表示第\(i\)项的系数

同时我们还可以发现\(E(\prod_{j=1}^mX_{j_i})=E(\prod_{i=1}^mX_i)\)，表示\(E\)中如果是个数相同的\(X_i\)相乘，那么他们的值总相等，于是这部分可以提取出来

那么\(E(S)=E(X_1X_2X_3)+(A_1+A_2+A_3)E(X_1X_2)+(A_1A_2+A_1A_3+A_2A_3)E(X_1)+A_1A_2A_3\)

? \(=S(0)E(X_1X_2X_3)+S(1)E(X_1X_2)+S(2)E(X_1)+S(3)E(1)\)

可以归纳总结出\(E(\prod_{i+1}^n(A_i+X_i))=\sum_{i=0}^nS(i)E(\prod_{j=1}^{n-i}X_j)\)

于是答案就很好求了，前面的\(S(i)\)可以先用\(CDQ\)分治或启发式合并利用\(NTT\)快速求\(O(n(logn)^2)\)

后面一部分考虑随机变量\(X_{i,j}\)表示第\(i\)个盒子在第\(j\)次操作是否被选中\((0/1)\)。则\(X_i=\sum_{j=1}^kX_{i,j}\)

\(E(\prod_{i=1}^nX_i)=E(\prod_{i=1}^n\sum_{j=1}^kX_{i,j})=\sum_{j_1,j_2,...,j_n}E(\prod_{i=1}^nX_{i,j_i})\)

后面的期望当且仅当\(X_{i,j_i}\)均为\(1\)的时候得1，否则均为0，而全部相等就每个\(i\)都被取到,取到\(i\)的概率为\(\frac{1}{N}\)，\(n\)个均被取到，则\((\frac{1}{N})^n\)。

因此 \(E(\prod_{i=1}^nX_i)=(\prod_{i=1}^n(K+1-i))(\frac{1}{N})^n=\frac{K!}{(K-n)!}(\frac{1}{N})^n=A_K^n(\frac{1}{N})^n\)

//#pragma GCC optimize(2)
#include 
#define ll long long
#define pb push_back
using namespace std;

const int N = 1100000;
const int p = 998244353, gg = 3, ig = 332738118, img = 86583718;//1004535809 
const int M=18e5;
const int mod=998244353;
template void rd(T &x)
{
    x = 0;
     int f = 1;
     char ch = getchar();
    while(ch < '0' || ch > '9') {if(ch == '-')f = -1;ch = getchar();}
    while(ch >= '0' && ch <= '9') {x = x * 10 + ch - '0';ch = getchar();}
    x *= f;
}

ll qpow(ll a, int b)
{
    ll res = 1;
    while(b) {
        if(b & 1) res = 1ll * res * a % mod;
        a = 1ll * a * a % mod;
        b >>= 1;
    }
    return res;
}

vectorA[N];
namespace Poly
{
    #define mul(x, y) (1ll * x * y >= mod ? 1ll * x * y % mod : 1ll * x * y)
    #define minus(x, y) (1ll * x - y < 0 ? 1ll * x - y + mod : 1ll * x - y)
    #define plus(x, y) (1ll * x + y >= mod ? 1ll * x + y - mod : 1ll * x + y)
    #define ck(x) (x >= mod ? x - mod : x)//取模运算太慢了

    typedef vector poly;
    const int G = 3;//根据具体的模数而定，原根可不一定不一样！！！
    //一般模数的原根为 2 3 5 7 10 6
    const int inv_G = qpow(G, mod - 2);
    int RR[N], deer[2][21][N], inv[N];

    void init(const int t) {//预处理出来NTT里需要的w和wn，砍掉了一个log的时间
        for(int p = 1; p <= t; ++ p) {
            int buf1 = qpow(G, (mod - 1) / (1 << p));
            int buf0 = qpow(inv_G, (mod - 1) / (1 << p));
            deer[0][p][0] = deer[1][p][0] = 1;
            for(int i = 1; i < (1 << p); ++ i) {
                deer[0][p][i] = 1ll * deer[0][p][i - 1] * buf0 % mod;//逆
                deer[1][p][i] = 1ll * deer[1][p][i - 1] * buf1 % mod;
            }
        }
        inv[1] = 1;
        for(int i = 2; i <= (1 << t); ++ i)
            inv[i] = 1ll * inv[mod % i] * (mod - mod / i) % mod;
    }

    int NTT_init(int n) {//快速数论变换预处理
        int limit = 1, L = 0;
        while(limit < n) limit <<= 1, L ++ ;
        for(int i = 0; i < limit; ++ i)
            RR[i] = (RR[i >> 1] >> 1) | ((i & 1) << (L - 1));
        return limit;
    }

    void NTT(poly &A, int type, int limit) {//快速数论变换
        A.resize(limit);
        for(int i = 0; i < limit; ++ i)
            if(i < RR[i])
                swap(A[i], A[RR[i]]);
        for(int mid = 2, j = 1; mid <= limit; mid <<= 1, ++ j) {
            int len = mid >> 1;
            for(int pos = 0; pos < limit; pos += mid) {
                int *wn = deer[type][j];
                for(int i = pos; i < pos + len; ++ i, ++ wn) {
                    int tmp = 1ll * (*wn) * A[i + len] % mod;
                    A[i + len] = ck(A[i] - tmp + mod);
                    A[i] = ck(A[i] + tmp);
                }
            }
        }
        if(type == 0) {
            for(int i = 0; i < limit; ++ i)
                A[i] = 1ll * A[i] * inv[limit] % mod;
        }
    }

    inline poly poly_mul(poly A, poly B) {//多项式乘法
        int deg = A.size() + B.size() - 1;
        int limit = NTT_init(deg);
        poly C(limit);
        NTT(A, 1, limit);
        NTT(B, 1, limit);
        for(int i = 0; i < limit; ++ i)
            C[i] = 1ll * A[i] * B[i] % mod;
        NTT(C, 0, limit);
        C.resize(deg);
        return C;
    }
    //多个多项式相乘CDQ或者利用优先队列启发式合并
    inline poly CDQ(int l,int r)
    {
      if(l==r)
      { 
        return A[l];
      }
      int mid=l+r>>1;
      poly L=CDQ(l,mid);
      poly R=CDQ(mid+1,r);
      return poly_mul(L,R);  
    }
}

using namespace Poly;
int n, k;
ll E[N];
int main()
{
    //freopen("in.in","r",stdin);
    init(20);//2^21 = 2,097,152,根据题目数据多项式项数的大小自由调整，注意大小需要跟deer数组同步(21+1=22)
    int n,k;
    rd(n),rd(k);
    for(int i=1;i<=n;i++)
    {
        int x;
        rd(x);
        A[i].pb(1);
        A[i].pb(x);
    }
    poly s=CDQ(1,n);
    int inv=qpow(n,mod-2);
    E[0]=1;
    for(int i=1;i<=n;i++)
        E[i]=E[i-1]*inv%mod*(k-i+1+mod)%mod;
    ll ans=0;
    for(int i=0;i<=n;i++)
        ans=(ans+s[i]*E[n-i]%mod)%mod;
    cout<