Solves the Riddler Classic from October 2, 2020: https://fivethirtyeight.com/features/can-you-eat-all-the-chocolates/

20201002-riddler.py

"""
Solves the Riddler Classic from October 2, 2020.
https://fivethirtyeight.com/features/can-you-eat-all-the-chocolates/
"""
from __future__ import annotations
from functools import lru_cache
from typing import List, NamedTuple, Optional, Tuple


class ChocolateBag(NamedTuple):
    """
    The ChocolateBag class is a container that specifies the number of chocolates left
    in the bag, and also provides useful methods for calculating the expected value of
    certain events, such as "the probability that the last chocolate pulled from the bag
    is a milk chocolate."

    To solve the Riddler Classic from October 2, 2020, you use this class with its
    default values, and can calculate the probability of drawing a milk chocolate last:

    >>> ChocolateBag(dark=8, milk=2).expected_value()
    0.5

    Parameters
    ----------
    dark : int, default 8, the number of dark chocolates remaining in the bag
    milk : int, default 2, the number of milk chocolates remaining in the bag
    last : Optional[str], default None, the type of chocolate last picked from the bag;
        if 'dark', then the last chocolate picked was a dark chocolate; if 'milk', then
        the last chocolate picked was a milk chocolate; if None, then no chocolate was
        picked last, which means we're starting fresh.
    """

    dark: int = 8
    milk: int = 2
    last: Optional[str] = None

    def step(self) -> List[Tuple[ChocolateBag, float]]:
        """
        Returns each possible future state we could reach from this position. The states
        are returned as a list of tuples where each tuple is (future_state, probability)

        Examples
        --------
        >>> steps = ChocolateBag(dark=8, milk=2, last=None).step()
        >>> steps[0]
        (ChocolateBag(dark=7, milk=2, last='dark'), 0.8)
        >>> steps[1]
        (ChocolateBag(dark=8, milk=1, last='milk'), 0.2)
        """
        p_dark = self.dark / (self.dark + self.milk)
        p_milk = self.milk / (self.dark + self.milk)

        if self.last == "dark":
            # the last chocolate we picked was dark, so we can either pick another one,
            # or we can reset to last=None if we pick a milk chocolate.
            return [
                (ChocolateBag(dark=self.dark - 1, milk=self.milk, last="dark"), p_dark),
                (ChocolateBag(dark=self.dark, milk=self.milk, last=None), p_milk),
            ]

        if self.last == "milk":
            # the last chocolate we picked was milk, so we can either pick another one,
            # or we can reset to last=None if we pick a dark chocolate.
            return [
                (ChocolateBag(dark=self.dark, milk=self.milk, last=None), p_dark),
                (ChocolateBag(dark=self.dark, milk=self.milk - 1, last="milk"), p_milk),
            ]

        # otherwise, we assume we're in a fresh state, and we can pick either chocolate
        return [
            (ChocolateBag(dark=self.dark - 1, milk=self.milk, last="dark"), p_dark),
            (ChocolateBag(dark=self.dark, milk=self.milk - 1, last="milk"), p_milk),
        ]

    @lru_cache()
    def expected_value(self):
        """
        Returns the probability of drawing a milk chocolate last from this ChocolateBag.
        This depends on the last chocolate drawn, and on those remaining in the bag.

        Examples
        --------
        >>> ChocolateBag(dark=0, milk=1, last=None).expected_value()
        1.0
        >>> ChocolateBag(dark=1, milk=1, last=None).expected_value()
        0.5
        >>> ChocolateBag(dark=1, milk=0, last=None).expected_value()
        0.0
        >>> ChocolateBag(dark=1, milk=1, last='dark').expected_value()
        0.75
        >>> ChocolateBag(dark=1, milk=1, last='milk').expected_value()
        0.25
        """
        # terminating states: we know the outcome when only one chocolate type is left
        if self.milk == 0:
            return 0.0
        if self.dark == 0:
            return 1.0

        # otherwise we branch into possible outcomes based on the chocolate we draw. The
        # branches are created by the step method, and we calculate the expected value.
        return sum(
            probability * state.expected_value() for state, probability in self.step()
        )


if __name__ == "__main__":
    import doctest

    doctest.testmod()