Prisoners in Hats

hats.py

import random
import itertools


def relative_complement(superset, subset):
    return [elem for elem in superset if elem not in subset]


def simulate():
    n = 1000

    all_hats = list(range(0, n + 1))

    prisoner_hats = list(all_hats)
    random.shuffle(prisoner_hats)
    prisoner_hats = prisoner_hats[:n]

    said_numbers = []
    saved = list(range(n))

    # Last prisoner in line
    # Get 2 unknown numbers
    unknown = relative_complement(all_hats, prisoner_hats[1:])
    mod_sum = sum(unknown) % (n + 1)
    saved[0] = (mod_sum == prisoner_hats[0])

    # Other prisoners
    for i in range(1, n):
        # Narrow down to 3 unknown numbers
        unknown = relative_complement(all_hats, prisoner_hats[(i + 1):])
        unknown = relative_complement(unknown, said_numbers)

        # Find which number is not part of mod_sum, and have prisoner "say" it
        for pair in itertools.combinations(unknown, 2):
            if (sum(pair) % (n + 1) == mod_sum):
                third_num = relative_complement(unknown, pair)[0]
                said_numbers.append(third_num)
                saved[i] = (third_num == prisoner_hats[i])
                break

    num_saved = len([e for e in saved if e is True])
    print("Saved: " + str(num_saved))
    print("Died: " + str(n - num_saved))
    print("Last prisoner in line died: " + str(not saved[0]))

simulate()