A python script to test strategies for the for the riddle described on TED-ed's youtube channel: www.youtube.com/watch?v=dh4nEuhZBgg

seven-planets-riddle

"""
A python script to test strategies for the for the riddle described on TED-ed's
youtube channel: www.youtube.com/watch?v=dh4nEuhZBgg

Your interstellar police squad has tracked a group of criminals to a cluster of
seven planets. Now you must apprehend them before their reinforcements arrive.
Of course, the fugitives won’t just stay put – they’ll try to dodge you by
moving from planet to planet. Can you devise a sequence for searching the
planets that’s guaranteed to catch them in ten warps or less?
Edwin F. Meyer shows how.

Planets are arranged like so:

  1       7
    2   6
      3
      4
      5

"""
import numpy as np

# Describe the possible transitions at each turn for each planet
PLANETS = {
    1: [2],
    2: [1, 3],
    3: [2, 4, 6],
    4: [3, 5],
    5: [4],
    6: [3, 7],
    7: [6]
}


def test(strategy, n=10000):
    """
    Test a given strategy's chance of catching the rogues by Monte-Carlo.
    """
    if len(strategy) != 10:
        raise Exception('Must have 10 planets')

    outcomes = []
    for run in range(n):
        rebel_position = np.random.choice(list(PLANETS))

        for turn, planet in enumerate(strategy):
            police_position = planet
            rebel_position = np.random.choice(list(PLANETS[rebel_position]))

            if police_position == rebel_position:
                outcomes.append({
                    'run': run,
                    'caught': True,
                    'turn': turn
                })
                break
        else:
            outcomes.append({
                'run': run,
                'caught': False,
                'turn': None
            })

    catches = sum((o['caught']) for o in outcomes)

    return catches, outcomes


if __name__ == '__main__':
    strategy = [2, 3, 4, 3, 6, 2, 3, 4, 3, 6]  # This is the winning strategy

    catches, outcomes = test(strategy)
    print('Caught rebels {:0.2f} percent of the time'.format(
        100 * catches / len(outcomes)
    ))