AtCoder Beginner Contest 231 G 問題 Balls in Boxes

問題

atcoder.jp

提出解答

atcoder.jp

問題の概要

$N$ 個の変数 $X_1, \dots, X_N$ があり, 最初は全て $0$ である. ここで, 以下を計 $K$ 回実行する.

$1$ 以上 $N$ 以下の整数 $i$ を等確率 (各回ごとで独立) で選び, $X_i$ に $1$ を加算する.

この操作後の $\displaystyle \prod_{i=1}^N (A_i+X_i)$ の期待値を求めよ.

制約

$1 \leq N \leq 1000$
$1 \leq K \leq 10^9$
$0 \leq A_i \leq 10^9$
入力はすべて整数

解法

$X_1=u_1, \dots, X_N=u_N$ となる場合を考える. 変数がこのようになる場合の数は

$\dfrac{K!}{u_1 ! \dots u_N!}$

である. 求めるべきは

$\begin{align} \displaystyle & \phantom{=} \dfrac{1}{N^K} \sum_{u_1+\dots+u_N=K} (A_1+u_1) \dots (A_N+u_N) \dfrac{K!}{u_1 ! \dots u_N!} \\ &=\dfrac{K!}{N^K} \sum_{u_1+\dots+u_N=K} \prod_{i=1}^N (A_i+u_i) \dfrac{1}{u_i !} \end{align}$

である.

ここで, この値を多項式の係数とみなす. 具体的には,

$\displaystyle \widetilde{F_i}=\sum_{n=0}^K (A_i+n) \dfrac{X^n}{n!}$

としたとき,

$\displaystyle \dfrac{K!}{N^K} \left(\left[ X^K \right] \prod_{i=1}^N \widetilde{F_i} \right)$

を求めることになる.

さて, このまま議論を進めても良いが, 今後の見通しを良くするために, $\widetilde{F_i}$ を

$\displaystyle F_i(X)=\sum_{n=0}^\infty (A_i+n) \dfrac{X^n}{n!}$

に置き換えて行う.

そうすると,

$\begin{align} F_i(X) &=\sum_{n=0}^\infty (A_i+n) \dfrac{X^n}{n!} \\ &=A_i \sum_{n=0}^\infty \dfrac{X^n}{n!}+\sum_{n=0}^\infty n \dfrac{X^n}{n!} \\ &=A_i \sum_{n=0}^\infty \dfrac{X^n}{n!}+\sum_{n=0}^\infty \dfrac{X^n}{(n-1)!} \\ &=A_i \sum_{n=0}^\infty \dfrac{X^n}{n!}+X \sum_{n=1}^\infty \dfrac{X^{n-1}}{(n-1)!} \\ &=A_i \sum_{n=0}^\infty \dfrac{X^n}{n!}+X \sum_{m=0}^\infty \dfrac{X^m}{m!} \\ &=(X+A_i) \sum_{m=0}^\infty \dfrac{X^m}{m!} \\ &=(X+A_i) e^X \end{align}$

である.

そして,

$\displaystyle \prod_{i=1}^N F_i(X)=\left(\prod_{i=1}^N (X+A_i) \right) \left(\prod_{i=1}^N e^X \right)=\left(\prod_{i=1}^N (X+A_i) \right) e^{NX}$

となる.

$\displaystyle H(X):=\prod_{i=1}^N (X+A_i)$ とおく. このとき,

$\begin{align} \left[X^K\right] (H e^{NX}) &=\sum_{l=0}^K \left(\left[X^l \right ] H \right) \left(\left[X^{K-l} \right]e^{NX} \right) \\ &=\sum_{l=0}^K \left(\left[X^l \right ] H \right) \dfrac{N^{K-l}}{(K-l)!} \\ &=N^K \sum_{l=0}^K \left(\left[X^l \right ] H \right) \dfrac{1}{N^l (K-l)!} \\ \end{align}$

となる.

ここで, $H$ は $N$ 次式なので, $l \gt N \Rightarrow \left[X^l \right] H=0$ である. よって,

$\displaystyle \left[X^K\right] (H e^{NX})=N^K \sum_{l=0}^{\min(N,K)} \left(\left[X^l \right ] H \right) \dfrac{1}{N^l (K-l)!} \\$

となる.

以上から, 求めるべき値は, $h_l:=\left[X^l \right ] H$ とすると,

$\begin{align} \dfrac{K!}{N^K} \left(\left[X^K\right] (H e^{NX}) \right) &=\dfrac{K!}{N^K} \left(N^K \sum_{l=0}^{K} \left(\left[X^l \right ] H \right) \dfrac{1}{N^l (K-l)!} \right) \\ &=K! \sum_{l=0}^{\min(N,K)} \dfrac{h_l}{N^l (K-l)!} \\ &=\sum_{l=0}^{\min(N,K)} \dfrac{h_l}{N^l} \dfrac{K!}{(K-l)!} \\ &=\sum_{l=0}^{\min(N,K)} \dfrac{h_l}{N^l} K(K-1)\cdots(K-l+1) \\ \end{align}$