prisoners.py

#!/usr/bin/env python3

import time
import random

from typing import Generator, Iterable


def extend_loop(boxes: tuple[int, ...], loop: list[int], pointer: int) -> tuple[int, ...]:
    """
    Find a loop in boxes.

    Args:
        boxes: E.g. (5, 1, 0, 2, 3, 4). Here there are 6 boxes. The first box contains a
            note which points to index 5. boxes[5] points to index 4, etc...
        loop: The current loop to extend. To start a new loop, pass an empty list.
        pointer: The index of the box to extend the loop from.

    Returns:
        A tuple of box indexes that form a loop.
    """
    next_box = boxes[pointer]

    if loop and next_box == loop[0]:  # Loop finished, so return it
        return tuple(loop)

    else:  # Loop not finished, so go to the next box
        loop.append(next_box)
        return extend_loop(boxes, loop, pointer=next_box)


def find_loops(boxes: tuple[int, ...]) -> Generator[tuple[int, ...], None, None]:
    """
    Generates loops present in a set of boxes.

    Args:
        boxes: See extend_loop.
    """
    not_in_loops = set(boxes)
    while not_in_loops:
        loop = extend_loop(boxes, loop=[], pointer=not_in_loops.pop())
        yield loop
        not_in_loops -= set(loop)


def loops_longer_than_x(n_boxes: int, x: int, n_samples: int) -> float:
    """
    Simulate n_boxes boxes, n_samples times. Return the proportion of times that the
    length of the longest loop is greater than x.

    Args:
        n_boxes: Number of boxes in each simulation.
        x: Length loop to exceed.
        n_samples: Number of simulations to run.
    """
    n = 0
    for _ in range(n_samples):
        boxes = list(range(n_boxes))
        random.shuffle(boxes)
        if any(len(loop) > x for loop in find_loops(boxes)):
            n += 1
    return n / n_samples


if __name__ == "__main__":
    t0 = time.time()
    n = loops_longer_than_x(n_boxes=100, x=50, n_samples=10_000)
    print("proportion: ", n, ", time: ", time.time() - t0)