E - Sum of Square of Sum Editorial by en_translator
For \(N=1\), the answer is obvious. Assume \(N>1\).
Instead of the original problem, consider the following:
Problem: find the expected score when \(K\) out of the \(N\) balls are chosen randomly.
The answer to the original problem is \(\binom{N}{K}\) times the answer to this one. Thus, it suffices to solve this problem to answer the original one.
Define the random variable \(X_i\) as \(A_i\) if ball \(i\) is chosen, and \(0\) if it is not.
Then the expected value is
\(\begin{aligned} E\left[\left(\sum_i X_i\right)^2\right] =\sum_i E[X_i^2] + \sum_{i\neq j}E[X_iX_j]. \end{aligned}\)
Here, by definition of the random variable, \(X_i^2\) equals \(A_i^2\) if ball \(i\) is chosen, and \(0\) otherwise, so \(E[X_i^2]=A_i^2\frac{K}{N}\). Also, \(X_iX_j\) equals \(A_iA_j\) if balls \(i\) and \(j\) were both chosen, and \(0\) otherwise, so \(E[X_iX_j]=A_iA_j\frac{K(K-1)}{N(N-1)}\).
Hence
\(\begin{aligned} E\left[\left(\sum_i X_i\right)^2\right] &=\sum_i E[X_i^2] + \sum_{i\neq j}E[X_iX_j]\\ &=\frac{K}{N}\sum_i A_i^2 + \frac{K(K-1)}{N(N-1)}\sum_{i\neq j}A_iA_j. \end{aligned}\)
By \(\sum_{i\neq j}A_i A_j = (\sum_i A_i)^2 - \sum_i A_i^2\), this can be evaluated with \(O(N)\) arithmetic operations.
Hence, the problem has been solved in about \(O(N+\log \mathrm{MOD})\) time.
posted:
last update: