Official

G - Balls in Boxes Editorial by en_translator


We denote the expected value of random variable \(X\) by \(E[X]\).

Let \(X_i\) be the random variable denoting how many times the box \(i\) was chosen in the \(K\) operations; what we want is \(E[\prod_i (A_i+X_i)]\). By the linearity of expected value, it can be decomposed into monomials. For example, when \(N=2\), we consider the decomposition \(E[(A_1+X_1)(A_2+X_2)]=E[A_1A_2]+E[A_1X_2]+E[X_1A_2]+E[X_1X_2]\).

Here, by the symmetry of operations, for \(n\) number of different indices \(i_1,\ldots,i_n\), it holds that \(E[\prod_{j=1}^{n} X_{i_j}]=E[\prod_{i=1}^{n} X_i]\). Therefore

\[E\left[\prod_{i=1}^{N} (A_i+X_i)\right]=\sum_{n=0}^{N}S_{N,n}(A_1,\ldots,A_N)E\left[\prod_{i=1}^{N-n}X_i\right],\]

where \(S_{N,n}\) denotes the elementary symmetric polynomial in \(N\) variables of degree \(n\).

Thus, if we can find \(S_{N,n}(A_1,\ldots,A_N)\) and \(E\left[\prod_{i=1}^{n}X_i\right]\) for every \(n\), we can find the answer for the original problem fast enough.

Computing the value of elementary symmetric polynomials

Let \(DP[i][j]=S_{i,j}(A_1,\ldots,A_i)\), then \(DP[i+1][j]=DP[i][j]+DP[i][j-1]A_{i+1}\), so we can easily find the values for every \(n\) in a total of \(O(N^2)\) time with DP (Dynamic Programming).
If you prefer interpreting them as polynomials, we may regard it as computing \(S_{N,n}(A_1,\ldots,A_N)=[X^n]\prod_{i=1}^{N}(1+A_iX)\).

Finding the expected values

Let \(X_{i,j}\) be a random variable that takes \(1\) if Box \(i\) was chosen in the \(j\)-th operation, and \(0\) otherwise. Since \(X_i=\sum_{j=1}^{K}X_{i,j}\), just as before we can decompose the expressions into monomials as

\[E\left[\prod_{i=1}^{n}X_i\right]=\sum_{j_1,\ldots,j_n}E\left[\prod_{i=1}^{n}X_{i,j_i}\right].\]

Also, by the semantics of the random variables, when \(i_1\neq i_2\) it holds that \(X_{i_1,j}X_{i_2,j}=0\), so

\[E\left[\prod_{i=1}^{n}X_{i,j_i}\right]= \begin{cases} \left(\frac{1}{N}\right)^n & \text{if }j_i\text{ are pairwise distinct }\\ 0 & \text{otherwise}\\ \end{cases}.\]

Therefore, we have

\[E\left[\prod_{i=1}^{n}X_i\right]=\left(\prod_{i=1}^{n}(K+1-i)\right)\left(\frac{1}{N}\right)^n,\]

which enables us to find the values for every \(n\) in the increasing order of \(n\) in a total of \(O(N)\) time.

Bounus: Solve it in a total of \(O(N(\log N)^2)\) time.

posted:
last update: