Solves the fivethirtyeight Riddler Express from March 22, 2019

riddler-baseball.py

# -*- coding: utf-8 -*-
"""
Solves the fivethirtyeight Riddler problem from March 22, 2019

Each inning is represnted by an object called "Inning", that is comprised of
strikes, outs, and runner positions, e.g. Inning(0, 1, (0,1,0)) for 0 strikes,
1 out, and a runner on second base.

There are 11 unique game actions that can occur. Each is represented by a 
function that takes an inning state as input and returns a tuple of the new 
inning state after the action and the number of runs that were scored as a 
result of the action.
"""
import numpy as np
from collections import namedtuple


# define a container for the current state of an inning
Inning = namedtuple('Inning', ['strikes', 'outs', 'runners'])

# all possible game actions, handled with functions below
__actions__ = [
    ('strike', 4),
    ('out_at_first', 4),
    ('fly_out', 3),
    ('single', 3),
    ('base_on_error', 1), 
    ('double', 1), 
    ('double_play', 1),
    ('foul_out', 1),
    ('homerun', 1),
    ('triple', 1),
    ('walk', 1)
]
ACTIONS = [a[0] for a in __actions__]
WEIGHTS = [float(a[1]) for a in __actions__]
WEIGHTS = [w / sum(WEIGHTS) for w in WEIGHTS]

def single_inning(verbose=True):
    """
    Simulate a single inning until three outs are recorded. Return the number
    of runs scored during the inning. If verbose=True, print a "play-by-play"
    summary of the inning for each action. If verbose=False, just return runs
    """
    state, score = Inning(0, 0, (0,0,0)), 0
    while state.outs < 3:
        action = np.random.choice(ACTIONS, p=WEIGHTS)
        state, runs = globals()[action](state)
        score += runs
        if verbose:
            print(f'{action:15}: {state}, {score} total runs scored')
    return score

def simulate_innings(trials):
    """Simulate a large number of innings and return the total runs scored"""
    return np.array([single_inning(verbose=False) for _ in range(trials)])

def base_on_error(inning):
    """
    Update the inning for a base on error
        * batter and runners advance one base each
    """
    if inning.runners == (0,0,0): runners, runs = (1,0,0), 0
    elif inning.runners == (1,0,0): runners, runs = (1,1,0), 0
    elif inning.runners == (0,1,0): runners, runs = (1,0,1), 0
    elif inning.runners == (0,0,1): runners, runs = (1,0,0), 1
    elif inning.runners == (1,1,0): runners, runs = (1,1,1), 0
    elif inning.runners == (1,0,1): runners, runs = (1,1,0), 1
    elif inning.runners == (0,1,1): runners, runs = (1,0,1), 1
    elif inning.runners == (1,1,1): runners, runs = (1,1,1), 1
    return Inning(0, inning.outs, runners), runs

def double(inning):
    """
    Update the inning for a double
        * all runners score; batter advances to second
        NOTE: in reality about 50% of runners on first score on a double
    """
    return Inning(0, inning.outs, (0,1,0)), sum(inning.runners)

def double_play(inning):
    """
    Update the inning for a double play
        * if the bases are empty this is treated like a single out; otherwise 
        it's assumed the two runners farthest advanced are picked off.
    """
    if inning.runners == (0,0,0):
        # only one out available
        runners = (0,0,0)
        outs = inning.outs + 1
    if inning.runners in [(1,0,0), (0,1,0), (0,0,1)]:
        runners = (0,0,0)
        outs = inning.outs + 2
    elif inning.runners in [(1,1,0), (1,0,1), (0,1,1)]:
        runners = (1,0,0)
        outs = inning.outs + 2
    elif inning.runners == (1,1,1):
        runners = (1,1,0)
        outs = inning.outs + 2
    else:
        runners = (0,0,0)
    if outs > 2:
        return  Inning(0, 3, (0,0,0)), 0
    else:
        return Inning(0, outs, runners), 0

def fly_out(inning):
    """
    Update the inning for a fly out
        * runner on third scores unless it's the third out
    """
    outs = inning.outs + 1

    if inning.runners == (0,0,1): runners, runs = (0,0,0), 1
    elif inning.runners == (1,0,1): runners, runs = (1,0,0), 1
    elif inning.runners == (0,1,1): runners, runs = (0,1,0), 1
    elif inning.runners == (1,1,1): runners, runs = (1,1,0), 1
    else: runners, runs = inning.runners, 0
    
    if outs > 2:
        return  Inning(0, 3, (0,0,0)), 0
    else:
        return Inning(0, outs, runners), runs

def foul_out(inning):
    """
    Update the inning for a foul out
        * runners stay where they are
    """
    if inning.outs > 1:
        return Inning(0, 3, (0,0,0)), 0
    else:
        return Inning(0, inning.outs + 1, inning.runners), 0

def out_at_first(inning):
    """
    Update the inning for an out at first
        * runners advance one base
    """
    outs = inning.outs + 1

    if inning.runners == (0,0,0): runners, runs = (0,0,0), 0
    elif inning.runners == (1,0,0): runners, runs = (0,1,0), 0
    elif inning.runners == (0,1,0): runners, runs = (0,0,1), 0
    elif inning.runners == (0,0,1): runners, runs = (0,0,0), 1
    elif inning.runners == (1,1,0): runners, runs = (0,1,1), 0
    elif inning.runners == (1,0,1): runners, runs = (0,1,0), 1
    elif inning.runners == (0,1,1): runners, runs = (0,0,1), 1
    elif inning.runners == (1,1,1): runners, runs = (0,1,1), 1
    
    if outs > 2:
        return  Inning(0, 3, (0,0,0)), 0
    else:
        return Inning(0, outs, runners), runs

def single(inning):
    """
    Update the inning for a single
        * runners advance two bases; batter advances to first
    """
    if inning.runners == (0,0,0): runners, runs = (1,0,0), 0
    elif inning.runners == (1,0,0): runners, runs = (1,0,1), 0
    elif inning.runners == (0,1,0): runners, runs = (1,0,0), 1
    elif inning.runners == (0,0,1): runners, runs = (1,0,0), 1
    elif inning.runners == (1,1,0): runners, runs = (1,0,1), 1
    elif inning.runners == (1,0,1): runners, runs = (1,0,1), 1
    elif inning.runners == (0,1,1): runners, runs = (1,0,0), 2
    elif inning.runners == (1,1,1): runners, runs = (1,0,1), 2
    return Inning(0, inning.outs, runners), runs

def strike(inning):
    """
    Update the inning state for a strike
    Returns a new Inning object and the number of runs scored
    """
    strikes = inning.strikes + 1
    if strikes > 2:
        strikes, outs = 0, inning.outs + 1
        if outs < 3: 
            return Inning(strikes, outs, inning.runners), 0
        else:     
            return Inning(0, 3, (0,0,0)), 0
    else:
        return Inning(strikes, inning.outs, inning.runners), 0

def triple(inning):
    """
    Update the inning for a triple
        * all runners score; batter advances to third
    """
    return Inning(0, inning.outs, (0,0,1)), sum(inning.runners)

def homerun(inning):
    """
    Update the inning for a triple
        * all runners and the batter score
    """
    return Inning(0, inning.outs, (0,0,0)), sum(inning.runners) + 1

def walk(inning):
    """
    Update the inning for a walk
        * batter advances to first; runners advance if forced
    """
    if inning.runners == (0,0,0): 
        runners, runs = (1,0,0), 0
    elif inning.runners in [(1,0,0), (0,1,0)]:
        runners, runs = (1,1,0), 0
    elif inning.runners == (0,0,1):
        runners, runs = (1,0,1), 0
    elif inning.runners in [(1,1,0), (1,0,1), (0,1,1)]:
        runners, runs = (1,1,1), 0
    elif inning.runners == (1,1,1):
        runners, runs = (1,1,1), 1
    return Inning(0, inning.outs, runners), runs


if __name__ == '__main__':
    trials = 100000
    results = simulate_innings(trials)
    print(f'Average score from {trials:,.0f} innings: {results.mean():,.2f}')