prisoners.py

import random
import logging

log = logging.getLogger(__name__)

MAX_NUM_PRISONERS = 12
NUM_PRISONERS = random.randrange(2, MAX_NUM_PRISONERS + 1)

def main():
    print(f'Number of prisoners: {NUM_PRISONERS}')
    prisoners = [Prisoner(i == 0) for i in range(NUM_PRISONERS)]

    switches = [False for _ in range(NUM_PRISONERS)]
    lights = {p: False for p in prisoners}

    while True:
        try:
            switches = [p.choice(lights[p]) for p in prisoners]
        except Answer as e:
            result = 'correct' if e.num == NUM_PRISONERS else 'incorrect'
            print(f'Prisoner answered: {e.num} ({result})')
            break

        log.debug(f'switches: {switches}')

        room_lights = switches[1:] + switches[:1]
        lights = {p: room_lights[i] for (i, p) in enumerate(prisoners)}

        random.shuffle(prisoners)

class Answer(Exception):

    def __init__(self, num):
        self.num = num

class Prisoner(object):

    def __init__(self, master):
        self.master = master
        self.choice_generator_instance = self.choice_generator()

    @staticmethod
    def flip_coin():
        return random.choice([False, True])

    def choice(self, light):
        self.light = light
        return next(self.choice_generator_instance)

    def choice_generator(self):

        def test_fill(active, iterations):
            for _ in range(iterations):
                yield active
                active = active or self.light
            for _ in range(2 ** iterations):
                yield active
                active = active and self.light
            self.fill_result = active

        # Determine an upper bound on the number of prisoners
        i = 1
        while True:
            log.debug(f'upper bound iteration: {i}')
            active = self.master
            yield from test_fill(active, i)
            if self.fill_result:
                break
            i += 1
        upper_bound = 2 ** i

        # Check if number of prisoners meeting a condition is nonzero
        def fill(active):
            for _ in range(upper_bound):
                yield active
                active = active or self.light
            for _ in range(upper_bound):
                yield active
                active = active and self.light
            self.fill_result = active
            log.debug(f'fill result: {self.fill_result}')

        # Assign the prisoners specific numbers
        if self.master:
            my_number = 0
        else:
            my_number = None

        num_assigned_prisoners = 1

        class TryAgain(Exception): pass

        while True:
            log.debug(f'my number: {my_number}')
            yield from fill(my_number is None)
            if not self.fill_result:
                # no non-assigned prisoners; we know the answer
                raise Answer(num_assigned_prisoners)
            while True:
                # in each iteration of this loop, we are attempting to assign a new number
                try:
                    sent_signal = self.flip_coin() if (my_number is not None) else False
                    yield sent_signal
                    received_signal = self.light
                    yield from fill(sent_signal)
                    if not self.fill_result:
                        log.info('no assigned prisoner sent signal; try again')
                        raise TryAgain()
                    num_sent_signal = 0

                    for assigned_prisoner_number in range(num_assigned_prisoners):
                        yield from fill(sent_signal and (assigned_prisoner_number == my_number))
                        if self.fill_result:
                            num_sent_signal += 1
                        if num_sent_signal == 2:
                            log.info('more than one assigned prisoner sent signal; try again')
                            raise TryAgain()

                    yield from fill(received_signal and (my_number is None))
                    if not self.fill_result:
                        log.info('the signal receiver is already assigned; try again')
                        raise TryAgain()
                    if received_signal:
                        assert my_number is None
                        my_number = num_assigned_prisoners
                        log.warning(f'assigning prisoner number: {my_number}')
                    num_assigned_prisoners += 1
                    break
                except TryAgain:
                    continue

if __name__ == '__main__':
    logging.basicConfig(level=logging.WARNING)
    main()