Riddler Classic: Can You Eat All The Chocolates?

dark_milk_chocolate.py

from collections import namedtuple, defaultdict
from fractions import Fraction

State = namedtuple("State", ["D", "M", "prev"])

init_state = State(D=8, M=2, prev=None)
state_space = defaultdict(Fraction)
state_space[init_state] = 1

while any(s.M + s.D > 0 for s in state_space):
    temp = defaultdict(Fraction)
    for s, p in state_space.items():
        if s.D + s.M > 0:
            if (pD := Fraction(s.D, s.D + s.M)) > 0:
                if s.prev in [None, 'D']:
                    temp[State(D=s.D-1, M=s.M, prev='D')] += p * pD
                else:
                    temp[State(D=s.D, M=s.M, prev=None)] += p * pD

            if (pM := Fraction(s.M, s.D + s.M)) > 0:
                if s.prev in [None, 'M']:
                    temp[State(D=s.D, M=s.M-1, prev='M')] += p * pM
                else:
                    temp[State(D=s.D, M=s.M, prev=None)] += p * pM
        else:
            temp[s] += p

    state_space = temp

s = State(D=0, M=0, prev='M')
print(f"Probability that last chocolate eaten was M: {state_space[s]}")