This Riddler puzzle about an interesting game involving picking colored balls out of a box. How long will the game last?
You play a game with four balls: One ball is red, one is blue, one is green and one is yellow. They are placed in a box. You draw a ball out of the box at random and note its color. Without replacing the first ball, you draw a second ball and then paint it to match the color of the first. Replace both balls, and repeat the process. The game ends when all four balls have become the same color. What is the expected number of turns to finish the game?
Extra credit: What if there are more balls and more colors?
Here is my solution to the first part (four balls):
[Show Solution]
Before I begin, I’d like to acknowledge Hector Pefo and Sawyer Tabony who also posted excellent solutions to this Riddler puzzle. We all arrived at the same answer (whew!) but our approaches differ slightly.
A natural way to think about this problem is that the game can exist in some number of states (the particular coloring of the balls in the box). At each step, at most one of the balls is recolored and we transition to another state. Such a collection of states and transition probabilities is known as a Markov Chain. Since we are only interested in how long it takes for all balls to be the same color, it’s not necessary to keep track of every possible state; we can aggregate states and simplify the problem.
Partition approach
One way to aggregate states is to count the number of balls of each different color without regard for the colors themselves. For example, 4 balls can be partitioned in 5 different ways:
\[
1+1+1+1,\qquad
2+1+1,\qquad
2+2,\qquad
3+1,\qquad
4
\]The partition “2+1+1”, for example, consists of cases where two of the balls are the same color and the other two balls are two other colors. Both the cases “red+red+green+blue” and “blue+blue+yellow+red” would fall into the category “2+1+1”. By using these five partitions as states in a Markov Chain, we can compute the transition probabilities to go from one state to the next. For example, the probability of transitioning from “2+1+1” to “3+1” is $\frac{1}{3}$ because in order for this transition to occur, we must first choose one of the two identically colored balls (probability $\frac{2}{4}$), then we must choose one of the other two balls out of the remaining three (probability $\frac{2}{3}$).
Here is what the complete Markov Chain looks like when all transition probabilities are filled in:
If we label the states in order, we can write the transition probabilities as a matrix that describes how states evolve.
\[
A = \begin{bmatrix}
0 & 0 & 0 & 0 & 0 \\
1 & \frac{1}{2} & 0 & 0 & 0 \\
0 & \frac{1}{6} & \frac{1}{3} & \frac{1}{4} & 0 \\
0 & \frac{1}{3} & \frac{2}{3} & \frac{1}{2} & 0 \\
0 & 0 & 0 & \frac{1}{4} & 0
\end{bmatrix}
\]For example, if $x$ is our current distribution, then $Ax$ is our distribution after one move. Note that the columns sum to $1$ because at each state, we must transition to some other state. The only exception is the last column, because the game stops once we reach the final state.
To compute the probability that the game ends after $k$ turns, we should find the probability that the game is in the final state after $k$ turns. Since we start in state $1$, the required quantity is $\left[ A^k \right]_{51}$. In other words, the $(5,1)$ component of $A^k$. We can compute this directly for each $k$ and we obtain the following distribution:
Note that the probability is zero for 1 or 2 turns because the game can’t end that quickly. It takes at least three turns (since there are four balls total) for all balls to be painted the same color. We can also compute the expected number of turns by evaluating the sum:
\[
A + 2A^2 + 3A^3 + \dots
= A(I-A)^{-2}
\]and then extracting the $(5,1)$ component. The result in this case turns out to be 9. When we compare the distribution to the mean in the figure above, we notice that the distribution has a heavy tail; although the mean is 9, the likeliest number of turns is actually 5.
Here is my solution to the general case with $N$ balls:
[Show Solution]
Although it’s possible to extend the partition used in the previous part to any other number of balls, in practice it is a challenge. One of the first obstacles is that the number of partitions as a function of $N$ is an intractable quantity. It is known as the partition function and is studied by number theorists for its deep connections to prime numbers. Partition functions were even mentioned in the recent movie about Ramanujan, “The Man Who Knew Infinity”. So we’re not taking that approach.
Counting Blue
We can still use Markov Chains, but we’ll have to come up with a different way of aggregating the states. One possibility is to define the states $k = 1,2,\dots,N$ which correspond to the number of Blue balls in the box. If Blue ends up winning, we will have started in state 1 and made our way eventually to state N without ever getting to state 0. The trick is to realize that no matter which color wins, the winning color has a Markov Chain like this one. All colors are indistinguishable so we might as well consider the case where Blue wins.
The next step is to compute transition probabilities. Here, we must be careful because we are conditioning on Blue winning. So we use Bayes’ Rule to compute the transitions. Given that there are $k$ Blue balls with $1 \le k \le N-1$, the probability that we transition $k\to k+1$ is:
\[
A_{k+1,k} = \frac{\mathbb{P}(k\to k+1 \text{, Blue wins})}{\mathbb{P}(\text{Blue wins})}
= \frac{\frac{k}{N}\cdot \frac{N-k}{N-1} \cdot \frac{k+1}{N}}{\frac{k}{N}} = \frac{(k+1)(N-k)}{N(N-1)}
\]Similarly, the probability that we transition $k\to k-1$ is:
\[
A_{k-1,k} = \frac{\mathbb{P}(k\to k-1 \text{, Blue wins})}{\mathbb{P}(\text{Blue wins})}
= \frac{\frac{N-k}{N}\cdot \frac{k}{N-1} \cdot \frac{k-1}{N}}{\frac{k}{N}} = \frac{(k-1)(N-k)}{N(N-1)}
\]The only other possibility is that we don’t transition at all, which is what happens the rest of the time. So $A_{k,k} = 1-A_{k+1,k}-A_{k-1,k}$. Simplifying the expressions, we obtain:
\[
A_{k+1,k} = \frac{(k+1)(N-k)}{N(N-1)},\quad
A_{k-1,k} = \frac{(k-1)(N-k)}{N(N-1)},\quad
A_{k,k} = 1-\frac{2k(N-k)}{N(N-1)}
\]As before, the $k^\text{th}$ column sums to 1. The only exception, as before, is the last column $A_{:,N}$, which is zero since the game ends once all colors are the same. The approach from here is identical to what we did when we used partitions. We can compute the probability of ending in $k$ turns by computing $\left[ A^k \right]_{N,1}$ and the expected number of turns is given by $\left[ A(I-A)^{-2} \right]_{N,1}$. Here are the distributions for $N=5$ and $N=8$.
As we can see, the distribution looks similar as we increase $N$. Curiously, the mean always appears to be $(N-1)^2$. In the next section, we will show that this is indeed the case in general!
Analytic expressions
Writing out the transition matrix explicitly, we have:
\[
A = \begin{bmatrix}
1-\frac{2(N-1)}{N(N-1)} & \frac{1(N-2)}{N(N-1)} & & & & 0 \\
\frac{2(N-1)}{N(N-1)} & 1-\frac{4(N-2)}{N(N-1)} & \frac{2(N-3)}{N(N-1)} & & & 0 \\
& \frac{3(N-2)}{N(N-1)} & 1-\frac{6(N-3)}{N(N-1)} & \ddots & & 0 \\
& & \frac{4(N-3)}{N(N-1)} & \ddots & \frac{(N-2)1}{N(N-1)} & 0 \\
& & & \ddots & 1-\frac{2(N-1)1}{N(N-1)} & 0 \\
& & & & \frac{N\cdot 1}{N(N-1)} & 0
\end{bmatrix}
\]If we define $a_k$ to be the expected number of turns until the game ends given that we are currently at state $k$, we can write the recursion:
\begin{align}
a_k &= 1 + A_{k+1,k} a_{k+1} + A_{k,k} a_k + A_{k-1,k} a_{k-1} \quad\text{for }k=1,\dots,N-1 \\
a_N &= 0
\end{align}Rearranging this expression, it looks roughly like $A^Ta+1=a$ (with the exception of the last row). Writing this as a single compact equation, we can factor things a bit and obtain a simpler equation:
\begin{multline}
\frac{1}{N(N-1)}
\begin{bmatrix}N-1 & & &\\ & \ddots & & \\ & & 2 & \\ & & & 1\end{bmatrix}
\begin{bmatrix}2 & -1 & & \\ -1 & 2 & \ddots & \\ & \ddots & \ddots & -1 \\ & & -1 & 2\end{bmatrix}\\
\times \begin{bmatrix}1 & & \\ & 2 & & \\ & & \ddots & \\ & & & N-1\end{bmatrix}
\begin{bmatrix}a_1 \\ a_2 \\ \vdots \\ a_{N-1}\end{bmatrix}
=
\begin{bmatrix}1 \\ 1 \\ \vdots \\ 1\end{bmatrix}
\end{multline}Or, simplified:
\[
\begin{bmatrix}2 & -1 & & \\ -1 & 2 & \ddots & \\ & \ddots & \ddots & -1 \\ & & -1 & 2\end{bmatrix}
\begin{bmatrix}a_1 \\ 2a_2 \\ \vdots \\ (N-1)a_{N-1}\end{bmatrix}
=
N(N-1)\begin{bmatrix}\frac{1}{N-1} \\ \frac{1}{N-2} \\ \vdots \\ 1\end{bmatrix}
\]By performing an LU decomposition, we may instead solve the following system of equations in sequence:
\begin{align}
\begin{bmatrix} 1 & & & & \\ -\frac{1}{2} & 1 & & & \\ & -\frac{2}{3} & 1 & & \\ & & \ddots & \ddots & \\ & & & -\frac{N-2}{N-1} & 1\end{bmatrix}
\begin{bmatrix} y_1 \\ y_2 \\ \vdots \\ y_{N-2} \\ y_{N-1} \end{bmatrix}
&= N(N-1)\begin{bmatrix}\frac{1}{N-1} \\ \frac{1}{N-2} \\ \vdots \\ \frac{1}{2} \\ 1\end{bmatrix}\\
\begin{bmatrix} 2 & -1 & & & \\ & \frac{3}{2} & -1 & & \\ & & \frac{4}{3} & \ddots & \\ & & & \ddots & -1 \\ & & & & \frac{N}{N-1} \end{bmatrix}
\begin{bmatrix} a_1 \\ 2a_2 \\ \vdots \\ (N-2)a_{N-2} \\ (N-1)a_{N-1} \end{bmatrix}
&= \begin{bmatrix}y_1 \\ y_2 \\ \vdots \\ y_{N-2} \\ y_{N-1}\end{bmatrix}
\end{align}The first system is easily solved for $y$ by forward-substitution. We first have $y_1=N$. Then, $y_2 = \frac{N(N-1)}{N-2}+\frac{1}{2}y_1$, and so on. The result we find is:
\[
y_k = N(N-1)\sum_{i=1}^k \frac{i}{(N-i)k}
\]The second system can be solved for $a$ by back-subsitution, starting from the final equation and working our way back up. Doing this, we find:
\[
a_k = \sum_{j=1}^{N-k} \frac{1}{k+j} y_{k+j-1}
\]We may now combine both expressions and obtain:
\[
a_k = N(N-1) \sum_{j=1}^{N-k} \frac{1}{k+j} \sum_{i=1}^{k+j-1} \frac{i}{(N-i)(k+j-1)}
\]This is a bit messy, but we only care about $k=1$ (the expected number of turns given that we start with only one blue ball and blue ends up winning). This sum can be simplified with a bit of work:
\begin{align}
a_1 &= N(N-1) \sum_{j=1}^{N-1} \frac{1}{j(j+1)} \sum_{i=1}^{j} \frac{i}{(N-i)} \\
&= N(N-1) \sum_{j=1}^{N-1} \left( \frac{1}{j}-\frac{1}{j+1}\right) \sum_{i=1}^{j} \frac{i}{(N-i)}
\end{align}Define $Q_j := \frac{1}{j+1}\sum_{i=1}^j \frac{i}{N-i}$ and the sum telescopes:
\begin{align}
a_1 &= N(N-1)\sum_{j=1}^{N-1} \left( \frac{1}{N-j} + Q_{j-1}-Q_j\right)\\
&= N(N-1)\left( \sum_{j=1}^{N-1}\frac{1}{N-j}-Q_{N-1} \right) \\
&= N(N-1)\left( \sum_{j=1}^{N-1}\frac{1}{N-j}-\frac{1}{N}\sum_{j=1}^{N-1}\frac{j}{N-j} \right) \\
&= N(N-1) \sum_{j=1}^{N-1}\frac{1}{N} \\
&= (N-1)^2
\end{align}And this is precisely what we were trying to show! The expected number of turns for a game with $N$ balls is $(N-1)^2$. A similar expansion can be done to compute $a_k$ for $k\ne 1$ but in this case there is only partial telescoping and we are left with
\[
a_k = (N-1)^2-\frac{N(N-1)}{k}\sum_{i=1}^k \frac{k-i}{N-i}
\]As $N$ gets large, the sum of reciprocals is well-approximated by a logarithm. If we denote $\alpha = k/N$ the fraction of balls that are blue, then the expected number of turns until all balls are blue is:
\[
\frac{a_k}{(N-1)^2} \approx \left(\frac{1-\alpha}{\alpha}\right)\log\left(\frac{1}{1-\alpha}\right)
\]Here is a scaled plot of this distribution of expected number of turns:
A conjecture?
I was intrigued by the distribution of the number of turns rather than just the expected value. As seen in the plots above, the distribution converges to something, but it’s not clear what. And there doesn’t appear to be an easy way to compute or approximate powers of $A$ analytically. I suspect the distribution tends to a log-normal distribution as $N$ gets large because when plotted on a log scale, it sure looks close to normal:
If you have any thoughts, leave a comment below!
Thanks for your solution.
I’m relieved to see 9 which I got doing essentially the same thing by the more primitive method of directly calculating the powers of A and summing multiples of the (5,1) entries, only for me they were the (1,4) entries since I used the transpose and I eliminated state (1,1,1,1). I like your trick of setting the (5,5) entry of A to zero. I set the corresponding entry to 1 and had to subtract before summing.
Since we know mean of E1 dstribution if we can calc mode (max prob) we can fit the lognormal. How to do that?
Great post! I also used Markov chains to solve the first part, but never knew about the trick where you set all the current absorbing state probabilities to 0.
One minor suggestion is to fully explain the $A_{k+1,k}$ equation:
\begin{align}
\mathbb{P}(k \to k+1, \textrm{Blue Wins}) &= \mathbb{P}(k \to k+1) \cdot \mathbb{P}(\textrm{Blue Wins} \mid k \to k+1) \\
&= \frac{N-k}{N} \frac{k}{N-1} \cdot \frac{k+1}{N}
\end{align}
Perhaps others saw this easily, but it took me a bit of time to see this result. When I increased N, I also found the distribution tends to the log-normal. I agree with your conjecture