December 18 2024
zack sent me a math puzzle. let’s do it out here.
Problem as stated: We have an urn with 9 red balls and 1 blue ball. When we pick a ball, it magically duplicates - we keep one and put the other back in the urn. We keep doing this until we have at least one of each color - what is the expected proportion of red balls?
This is horribly stated. Let’s restate it.
We have a biased coin that flips heads with $P[H] = 0.9$ and tails with $P[T] = 0.1$. We flip the coin until we get at least one of each. What is the expected proportion of heads?
Intuitively, it may feel like the expectation is just $0.9$, because $0.9$ is the proportion of heads. It turns out it’s slightly less due to our best friend Jensen’s inequality.
Let’s split it up into two cases, based off of the first flip. It’s either heads or tails, with different expectations for each.
Let $X$ denote the proportion of heads, and $F$ be the event that we flip heads first. Then, we have $E[X] = E[X|F]P[F] + E[X|F^c]P[F^c].$
Let’s investigate the tails case, or $F^c$. We first see that $P[F^c] = 0.1$, so it remains to calculate $E[X | F^c]$. Notice that now, our stopping condition is getting a single heads. So we can model this as a geometric random variable.
Concretely, let $Y$ be the number of flips until we get a heads. Then, $Y \sim \text{Geom}(0.9),$ and we see that after $Y$ flips we now have 1 head and $Y$ tails (including the first flip), so the proportion of heads is $\frac{1}{Y+1}.$ In expectations, this is written as $E[X | F^c] = E\left[\frac{1}{Y+1}\right].$ If $E[1/X] = 1/E[X]$, then it turns out that 0.9 would probably be the expected value here. But it’s not, so we have to do more work.
Let’s calculate $E\left[\frac{1}{Y+1}\right]$ explicitly, with $p = 0.9$.
$$ E\left[\frac{1}{Y+1}\right] = \sum_{k=1}^\infty \frac{p \cdot (1-p)^{k–1}}{k+1} = \frac{p}{(1-p)^2} \sum_{k=1}^\infty \frac{(1-p)^{k+1}}{k+1}. $$
Let $b = 1-p$. We are now interested in finding $\sum_{k=1}^\infty \frac{b^{k+1}}{k+1}$. Reindexing, this gives us $\sum_{k=2}^\infty \frac{b^k}{k}$. One might recall an integration trick from finding the expected value of the geometric distribution. It turns out we can use the same integration trick to do our job.
Let’s consider $\sum_{k=1}^\infty b^k$. This is a geometric series, and since $0 < b < 1$ it converges. Thus, we have that
$$ \int_0^b \sum_{k=1}^\infty x^k dx = \sum_{k=1}^\infty \int_0^b x^k dx = \sum_{k=1}^\infty \frac{b^{k+1}}{k+1} = \sum_{k=2}^\infty \frac{b^k}{k}. $$
But $\sum_{k=1}^\infty b^k = \frac{b}{1-b}$, so we see that $$ \sum_{k=2}^\infty \frac{b^k}{k} = \int_0^b \frac{x}{1-x} dx = -b - \ln(1-b). $$
Let’s relate this back to our original sum.
$$ E\left[\frac{1}{Y+1}\right] = \frac{p}{(1-p)^2} \sum_{k=2}^\infty \frac{(1-p)^k}{k} = \frac{p}{(1-p)^2} \cdot (-(1-p) - \ln(1 - (1-p)) $$
To simplify this a little bit, let $q = 1-p$. We then have $$E\left[\frac{1}{Y+1}\right] = \frac{-(pq + p\ln(p))}{q^2}$$.
Thus, if we plug in $p = 0.9, q = 0.1,$ we’ll end up with the expected value of proportion of heads, given that we flip a tails on the first flip. This turns out to be $\approx 0.482$, which makes sense - in all likelyhood the next one is red, giving us a proportion of $1:1$, but this isnt always the case so it should be slightly lower.
One must wonder - can we do the same thing on the other side? The answer is yes. Let us think about the proportion of heads, given that we flip a heads first. We are now counting the number of flips $Z$ until we get a tails, which turns out to be another geometric random variable with $p = 0.1$. We will end up with $Z$ heads (including first flip) and $1$ tail, so the total proportion is $Z/(Z+1)$. Thus, we’re interested in $E[Z/(Z+1)]$ which isn’t quite the same form. However, notice that $Z/(Z+1) = 1 – 1/(Z+1)$, so all is good. Let’s put everything together.
$$ E[X] = E[X | F] P[F] + E[X | F^c] P[F^c] $$ $$ = \boxed{10 \left(-\frac{9}{100} + \frac{9}{10} \ln\left(\frac{10}{9}\right)\right) + \frac{9}{10} \left(1 - \frac{100}{81} \left(-\frac{9}{100} + \frac{\ln(10)}{10}\right)\right) \approx 0.7924} $$
As a final sanity check, we can write a little simluation to do this for us
import random
def expected_proportion(p):
H = 0
T = 0
while H == 0 or T == 0:
if random.random() < p:
H += 1
else:
T += 1
return H / (H + T)
tries = 1000000
p = 0.9
random.seed(0)
prop = 0
for i in range(tries):
prop += expected_proportion(p)
print(prop / tries)Running this gives us $0.7924890799861707$, which is pretty close (up to 4 sig figs!) to our expected value.