Detailed mathematical approach

Problem Summary

The implementations use an equivalent tracked-ball formulation of the urn process. At stage \(t\), let

$m=n-t+2,\qquad P_t=km.$

The active population has \(P_t\) balls, and \(X_t\) of them are the balls whose future influence must still be tracked. We draw \(2k\) balls uniformly without replacement, and \(B_t\) denotes how many tracked balls are removed in that draw. The quantity to evaluate is

$S(n,k)=\mathbb{E}\left[\sum_{t=1}^{n} B_t^2\right].$

The direct state distribution becomes too large when \(n\) is large, so the key idea is to evolve only the first two moments of \(X_t\). The sample checkpoint used by the implementation is \(S(2,2)=9.6\).

Mathematical Approach

The full distribution of the process is unnecessary once we notice that both the step contribution and the next-state update depend only on the first two moments of the current tracked-ball count.

Step 1: Write the Process in State Form

Before the draw at stage \(t\), there are \(X_t\) tracked balls inside a population of size \(P_t=km\), where \(m=n-t+2\). Conditional on \(X_t=x\), the number drawn from the tracked set is hypergeometric:

$\Pr(B_t=b\mid X_t=x)=\frac{\binom{x}{b}\binom{km-x}{2k-b}}{\binom{km}{2k}}.$

After the draw, the next stage keeps the undrawn tracked balls and adds one fresh block of \(k\) tracked balls, so for \(t<n\),

$X_{t+1}=X_t+k-B_t,$

with initial value

$X_1=k.$

This identity is the backbone of the whole method: once we know the moments of \(X_t\), we can compute both the expected contribution of stage \(t\) and the moments of \(X_{t+1}\).

Step 2: Compute the One-Step Contribution

For a hypergeometric random variable, the first factorial moments are standard:

$\mathbb{E}[B_t\mid X_t=x]=\frac{2k}{km}x=\frac{2}{m}x,$

$\mathbb{E}[B_t(B_t-1)\mid X_t=x]=\frac{2k(2k-1)}{km(km-1)}x(x-1)=\frac{2(2k-1)}{m(km-1)}x(x-1).$

Define

$\alpha_m=\frac{2}{m},\qquad \gamma_m=\frac{2(2k-1)}{m(km-1)}.$

Because \(B_t^2=B_t(B_t-1)+B_t\), we get

$\mathbb{E}[B_t^2\mid X_t=x]=\gamma_m x^2+(\alpha_m-\gamma_m)x.$

If we denote

$\mu_t=\mathbb{E}[X_t],\qquad \nu_t=\mathbb{E}[X_t^2],$

then the expected contribution of stage \(t\) is

$\mathbb{E}[B_t^2]=\gamma_m\nu_t+(\alpha_m-\gamma_m)\mu_t.$

Step 3: Derive the First-Moment Recurrence

Starting from \(X_{t+1}=X_t+k-B_t\), condition on \(X_t=x\):

$\mathbb{E}[X_{t+1}\mid X_t=x]=x+k-\mathbb{E}[B_t\mid X_t=x]=x+k-\frac{2}{m}x.$

Hence

$\mathbb{E}[X_{t+1}\mid X_t=x]=\frac{m-2}{m}x+k.$

Define

$\rho_m=\frac{m-2}{m}.$

Taking expectations gives the linear recurrence

$\mu_{t+1}=\rho_m\mu_t+k,$

with initial value \(\mu_1=k\). This already shows why the process is manageable: the next mean depends only on the current mean.

Step 4: Derive the Second-Moment Recurrence

Now square the state transition:

$X_{t+1}^2=(X_t+k-B_t)^2.$

Conditioning on \(X_t=x\) and inserting the formulas from Step 2 yields

$\mathbb{E}[X_{t+1}^2\mid X_t=x]=\mathbb{E}[(x+k-B_t)^2\mid X_t=x].$

After collecting the coefficients of \(x^2\), \(x\), and the constant term, the result is

$\mathbb{E}[X_{t+1}^2\mid X_t=x]=\sigma_m x^2+\left(2k\rho_m+\rho_m-\sigma_m\right)x+k^2,$

where

$\sigma_m=\frac{(m-2)(k(m-2)-1)}{m(km-1)}.$

Taking expectations gives

$\nu_{t+1}=\sigma_m\nu_t+\left(2k\rho_m+\rho_m-\sigma_m\right)\mu_t+k^2,$

with initial value \(\nu_1=k^2\). Therefore the whole process closes on the pair \((\mu_t,\nu_t)\), and no larger state object is required.

Worked Example: \(n=2,\ k=2\)

At the first stage, \(m=3\), so the population has \(6\) balls and \(X_1=2\). We draw \(4\) balls. The hypergeometric distribution is

$\Pr(B_1=0)=\frac{\binom{2}{0}\binom{4}{4}}{\binom{6}{4}}=\frac{1}{15},$

$\Pr(B_1=1)=\frac{\binom{2}{1}\binom{4}{3}}{\binom{6}{4}}=\frac{8}{15},$

$\Pr(B_1=2)=\frac{\binom{2}{2}\binom{4}{2}}{\binom{6}{4}}=\frac{6}{15}.$

Therefore

$\mathbb{E}[B_1^2]=0^2\cdot\frac{1}{15}+1^2\cdot\frac{8}{15}+2^2\cdot\frac{6}{15}=\frac{32}{15}.$

Next, \(X_2=4-B_1\), so \(X_2\) takes values \(4,3,2\) with the same probabilities, and

$\mathbb{E}[X_2^2]=4^2\cdot\frac{1}{15}+3^2\cdot\frac{8}{15}+2^2\cdot\frac{6}{15}=\frac{112}{15}.$

At the second stage, \(m=2\), so the draw size \(2k\) equals the whole remaining population. Hence \(B_2=X_2\), and

$\mathbb{E}[B_2^2]=\mathbb{E}[X_2^2]=\frac{112}{15}.$

The final expectation is

$S(2,2)=\frac{32}{15}+\frac{112}{15}=\frac{144}{15}=\frac{48}{5}=9.6,$

which matches the checkpoint embedded in the implementation.

How the Code Works

The C++, Python, and Java implementations all follow the same recurrence. They store only the current first moment \(\mu_t\), the current second moment \(\nu_t\), and the running total of \(\mathbb{E}[B_t^2]\). For each stage they compute \(m=n-t+2\), evaluate the contribution formula

$\gamma_m\nu_t+(\alpha_m-\gamma_m)\mu_t,$

add it to the total, and then update \(\mu_t\) and \(\nu_t\) with the linear recurrences above unless the last stage has been reached.

The Python and Java versions use high-precision decimal arithmetic, while the C++ version uses extended floating-point arithmetic. The C++ implementation also contains tiny-case self-checks that propagate the full hypergeometric distribution and compare it against the moment recurrence. After the loop, the final expectation is rounded to the nearest integer before being printed.

Complexity Analysis

The production algorithm performs one constant-size update per stage, so it runs in \(O(n)\) time and uses \(O(1)\) memory. This is the decisive improvement over a full distributional dynamic program, whose state space grows with the possible tracked-ball counts. The small validation routine is much more expensive because it keeps the whole probability distribution, but it is used only on tiny test cases.

Footnotes and References

1. Problem page: https://projecteuler.net/problem=783
2. Hypergeometric distribution: Wikipedia — Hypergeometric distribution
3. Urn problem: Wikipedia — Urn problem
4. Sampling without replacement: Wikipedia — Simple random sample
5. Law of total expectation: Wikipedia — Law of total expectation