Solves the fivethirtyeight Riddler from Oct 2, 2020.

chocolates.py

from typing import Optional


def model(dark: int, milk: int, last: Optional[str] = None) -> float:
    """
    Solves the Riddler Classic using dynamic programming. We start with a given number
    of chocolates in the bag, and we draw one at a time. We track the last chocolate we
    drew, and if we draw the same one, we eat it. Otherwise, we put it back in the bag
    and reset the "last" value to None.

    For example, `model(8, 2, None)` will return the probability of drawing a milk
    chocolate last if we start with a bag of 8 dark and 2 milk chocolates.

    Parameters
    ----------
    dark : int, the number of dark chocolates
    milk : int, the number of milk chocolates
    last : Optional[str], the type of chocolate we drew last, either "dark", "milk", or
        None (for example, if we're starting fresh)

    Returns
    -------
    expected_value : float, the probability of drawing a milk chocolate last, given our
        strategy of choosing chocolates.
    """
    if dark == 0:
        return 1.0
    if milk == 0:
        return 0.0
    p = dark / (dark + milk)
    if last == "dark":
        return p * model(dark - 1, milk, "dark") + (1 - p) * model(dark, milk, None)
    if last == "milk":
        return p * model(dark, milk, None) + (1 - p) * model(dark, milk - 1, "milk")
    return p * model(dark - 1, milk, "dark") + (1 - p) * model(dark, milk - 1, "milk")