hats.py

"""Prisoners in hats problem

N prisoners are notified that they will be playing a game tomorrow morning.
First they will stand in a circle and close their eyes.
Hats of one of N colors will be put on their heads (hat colors may repeat).
They will open their eyes and see the colors of everyone else's hat.
They will close their eyes and simultaneously guess what their own hat color is.

If any one of them guesses correctly, they all go free.
Otherwise they all die.

The prisoners may work together to form a strategy beforehand.

"""

###########################
# WARNING: SPOILERS AHEAD #
###########################

N = 7

def make_guess(N, prisoner_id, others_colors):
    return (prisoner_id - sum(others_colors)) % N

all_numbers = lambda num_digs, base=2: (tuple(x/(base ** i) % base for i in reversed(xrange(num_digs))) for x in range(num_digs ** base))
print all(any(make_guess(N, i, [x for j, x in enumerate(p) if i != j]) == p[i] for i in range(len(p))) for p in all_numbers(N, N))