riddler-dnd.py

# Solution to https://fivethirtyeight.com/features/can-you-find-the-best-dungeons-dragons-strategy/

# Results
# expectation [10.50030281  9.83376316 11.16661513]
# best_option [0 2 2 2 2 2 2 2 2 2 2 2 2 0 0 0 0 0 0 0]
# worst_option [0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1]

import numpy as np

SAMPLES = int(1e9)
BATCH_SIZE = int(1e6)

roll = lambda: np.random.randint(20, size=BATCH_SIZE, dtype='uint8')
count = lambda x: np.unique(x, return_counts=True)[1]
get_counts = lambda: np.vstack((
    count(roll()),
    count(np.maximum(np.minimum(roll(), roll()), np.minimum(roll(), roll()))),
    count(np.minimum(np.maximum(roll(), roll()), np.maximum(roll(), roll())))
))

counts = get_counts()
for i in range(SAMPLES // BATCH_SIZE - 1):
    counts += get_counts()
buckets = np.tile(np.arange(20) + 1, (3, 1))
expectation = np.average(buckets, weights=counts, axis=1)
sums = np.flip(np.cumsum(np.flip(counts), axis=1))
best_option = np.argmax(sums, axis=0)
worst_option = np.argmin(sums, axis=0)

print('expectation {}\nbest_option {}\nworst_option {}'.format(
    expectation, best_option, worst_option))