ABC231: G - Balls in Boxes

問題概要

1 から N までの番号がついた N 個の箱があり，箱 i (1\leq i \leq N) に入っているボールの数は A_i である。 N 個の数から等確率で 1 つ選んでボールを 1 足すという操作を K 回 繰り返す。箱 i (1\leq i \leq N) に操作後に入っているボールの数を B_i として \prod_{i=1}^{N} B_i の期待値はいくらか。

(問題文は https://atcoder.jp/contests/abc231/tasks/abc231_g に記載されている。)

制約

  • 1 \leq N \leq 1000
  • 1 \leq K \leq 10^9
  • 1 \leq A_i \leq 10^9 (1\leq i \leq N)

アイディア

あとで答えを N^K で割ることにして，すべての操作における \prod_{i=1}^{N} B_i の和を求めることを考える。

まず重要なのは，各操作における \prod_{i=1}^{N} B_i の総和を求めることは，各操作を終えた後に N 個の箱それぞれからボールを 1 つずつ取り出す方法を数えるのと同じである。そこで，操作前に箱に入っていたボールを赤色に，操作で入れるボールを白色に塗ることにする。このとき，最後に選んだボールのうち赤色に塗られたものを取り出した箱の添字の集合を R，白色に塗られたものを取り出した箱の添字の集合を Wとする。このとき R \cap W = \varnothingかつ R\cup W = \{ i\in \mathbb{Z}\,:\, (1\leq i \leq N)\}である。また，箱 iに入っている白色のボールの個数を D_iとする。 D_iに関して以下が成り立つ。

  • D_i \geq 0 (1\leq i \leq N)
  • D_i \geq 1 (i \in W)
  • \sum_{i=1}^{N} D_i = K

この状況に限定した場合における「操作後に N 個の箱それぞれからボールを 1 つずつ取り出す方法」は \left(\prod_{i\in W} D_i\right) \left( \prod_{i\in R} A_i \right) である。また，このような状況になる場合の数は \dfrac{K!}{\prod_{i\in W} D_i!} である。ところで，これらをかけ合わせると

\displaystyle \left(\prod_{i\in W} D_i\right) \left( \prod_{i\in R} A_i \right) \dfrac{K!}{\prod_{i=1}^{N} D_i!}

\displaystyle = \left( \prod_{i\in R} A_i \right) \dfrac{K! \prod_{i\in W} D_i}{\prod_{i=1}^{N} D_i!}

\displaystyle = \left( \prod_{i\in R} A_i \right) \dfrac{K!}{\left(\prod_{i\in W} (D_i-1)!\right) \left(\prod_{i\in R} D_i!\right)}

となる。ここで， Wを固定した場合のすべての適切な (D_1, D_2, \cdots, D_N)の組み合わせについて \dfrac{K!}{\left(\prod_{i\in W} (D_i-1)!\right) \left(\prod_{i\in R} D_i!\right)}の総和を s(W)とする。 s(W)が |W| *1にのみ依存することを示す。

今， C_i(1\leq i \leq N) を

  • C_i = D_i (i \in R)
  • C_i = D_i-1 (i \in W)

とおく。すると， \dfrac{K!}{\left(\prod_{i\in W} (D_i-1)!\right) \left(\prod_{i\in R} D_i!\right)}は \dfrac{K!}{\prod_{i=1}^{N} C_i!}と書き直せる。また， \sum_{i=1}^{N} C_i = ( \sum_{i\in R} D_i ) + ( \sum_{i\in W} (D_i-1) ) = \sum_{i=1}^{N} D_i - |W| = K-|W|であり， C_i \geq 0 (1\leq i \leq N)である。つまり， s(R)は \sum_{i=1}^{N} C_i = K-|W|， C_i \geq 0 (1\leq i \leq N)を満たすすべての (C_1, C_2, \cdots, C_N)の組み合わせについて \dfrac{1}{\prod_{i=1}^{N} C_i!}を求めその総和をとったものに s(W)をかけたものとなる。逆に， \sum_{i=1}^{N} C_i = K-|W|， C_i \geq 0 (1\leq i \leq N)を満たすすべての (C_1, C_2, \cdots, C_N)の組み合わせについて \dfrac{1}{\prod_{i=1}^{N} C_i!}の総和を取ると \dfrac{s(W)}{K!}となる。ところで，

\displaystyle \exp(x) = \sum_{i=0}^{\infty} \dfrac{x^i}{i!}

より

\displaystyle \exp(Nx) = \left(\sum_{i=0}^{\infty} \dfrac{x^i}{i!}\right)^N = \left(\sum_{i=0}^{\infty} \dfrac{x^i}{i!}\right) \left(\sum_{i=0}^{\infty} \dfrac{x^i}{i!}\right)\cdots \left(\sum_{i=0}^{\infty} \dfrac{x^i}{i!}\right)

となるが， exp(Nx)の x^{K-|W|}の係数がちょうど \dfrac{s(W)}{K!}となっている *2。 \exp(Nx)の x^{K-|W|}の係数を求めるには \exp(Nx)を {K-|W|}回微分してから x=0を代入して (K-|W|)!で割ればよく，実際にやると \dfrac{N^{K-|W|}}{(K-|W|)!}となる。よって \dfrac{s(W)}{K!} = \dfrac{N^{K-|W|}}{(K-|W|)!}なので s(W) = N^{K-|W|} \prod_{i=0}^{|W|-1} (K-i)となり， |W|にのみ依存することが示された。

これによって

\displaystyle \left( \prod_{i\in R} A_i \right) \dfrac{K!}{\left(\prod_{i\in W} (D_i-1)!\right) \left(\prod_{i\in R} D_i!\right)}

は

\displaystyle \left( \prod_{i\in R} A_i \right) N^{K-|W|} \prod_{i=0}^{|W|-1} (K-i)

と書き換えられる。結局求めたいものはすべての Wについてのこの総和である。ここでは |W|が同じものをまとめて計算することを考えると， |W|が同じであるすべての Wについて \prod_{i\in R} A_iの総和が求まればよく，これは簡単な O(N^2)のDPで可能である。

実装

まず dp(i,j)を「 i個目までの数から j個選ぶ際の選んだ数の積」の和としてこれを求める。初期状態は dp(0,0) = 1で，遷移は dp(i,j) = dp(i-1,j)+dp(i-1,j-1)\cdot A_i (ただし dp(\cdot, -1)=0とする)となる。

また， calc(i) := \prod_{j=0}^{i-1} (K-j)も求めておく。 calc(0) = 1, calc(i+1) = calc(i)\cdot(K-i)である。

そして \sum_{i=0}^{N} dp(N,N-i)\cdot calc(i)\cdot N^{K-i} を求め，最後に N^Kで割る。あるいは最初から \sum_{i=0}^{N} dp(N,N-i)\cdot calc(i)\cdot N^{-i} を求めてもよいかもしれない。

コード

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

const long long int MOD = 998244353LL;

long long int dp[1001][1001],calc[1001],A[1000];

long long int pow_binary(long long int x, long long int y){
    long long int ret = 1LL;
    while(y){
        if(y%2){
            ret *= x; ret %= MOD;
        }
        x *= x; x %= MOD;
        y /= 2LL;
    }
    return ret;
}

int main(void){
    int N,i,j;
    long long int K,perK,V=1LL,ANS=0LL;
    cin >> N >> K;
    perK = pow_binary(N,MOD-2);
    for(i=0; i<N; ++i){
        cin >> A[i]; A[i] %= MOD;
    }
    dp[0][0] = 1LL;
    calc[0] = 1;
    for(i=0; i<N; ++i){
        calc[i+1] = (calc[i]*(K-i))%MOD;
        for(j=0; j<=i; ++j){
            dp[i+1][j] += dp[i][j]; dp[i+1][j] %= MOD;
            dp[i+1][j+1] += (dp[i][j]*A[i])%MOD; dp[i+1][j+1] %= MOD;
        }
    }
    for(i=0; i<=N; ++i){
        calc[i] *= dp[N][N-i]; calc[i] %= MOD;
        calc[i] *= V; calc[i] %= MOD;
        ANS += calc[i]; ANS %= MOD;
        V *= perK; V %= MOD;
    }
    cout << ANS << endl;
    return 0;
}

(提出したものはこちら)

感想

結構悩んで面白かった。分かる人からしたら基礎中の基礎なんだろうけど。

*1: Wのサイズで，つまり白色に塗られた石の数。

*2: \left(\sum_{i=0}^{\infty} \dfrac{x^i}{i!}\right) \left(\sum_{i=0}^{\infty} \dfrac{x^i}{i!}\right)\cdots \left(\sum_{i=0}^{\infty} \dfrac{x^i}{i!}\right) を展開しようとしてみると分かるはず。