Code to generate solutions to the prison chessboard puzzle

main.py

from collections import deque

def solve(bits: int):
    # build a directed graph of node->child and also record mappings of node->parent
    digraph = {}
    parent_mappings = {}
    for n in list(range(1 << bits)):
        connections = _get_connected_numbers(n, bits)
        for c in connections:
            if c not in digraph or n not in digraph[c]:
                if n not in digraph:
                    digraph[n] = []
                digraph[n].append(c)

                if c not in parent_mappings:
                    parent_mappings[c] = []
                parent_mappings[c].append(n)

    # determine the colors
    colors = _color_nodes(digraph, parent_mappings, bits)

    return (digraph, colors)

# determine the colors via breadth-first search
def _color_nodes(
    graph: dict[int, list[int]], parent_mappings: dict[int, list[int]], bits: int
):
    colors: dict[int, int] = {}

    nodes = deque()
    first_node = list(graph.keys())[0]
    nodes.append(first_node)

    def process_node(node: int):
        # get the parents of this node:
        parents = parent_mappings[node] if node in parent_mappings else []

        # start with all colors as "possible"
        possible_colors = set(range(bits))

        for p in parents:
            # for each of the children of each parent node, if that node has a color set, remove it from the set of possible colors
            children_of_parent = graph[p] if p in graph else []
            for n in children_of_parent:
                if n in colors and colors[n] in possible_colors:
                    possible_colors.remove(colors[n])

            # for each of the parents of each parent node, if that node has a color set, remove it from the set of possible colors
            parents_of_parent = parent_mappings[p] if p in parent_mappings else []
            for n in parents_of_parent:
                if n in colors and colors[n] in possible_colors:
                    possible_colors.remove(colors[n])

        # set color to the lowest of the remaining possible colors
        colors[node] = min(possible_colors)

        # return if no children:
        if node not in graph:
            return

        # continue by traversing the children (if they haven't been colored yet and aren't in the queue yet)
        children = [c for c in graph[node] if c not in colors and c not in nodes]
        for child in children:
            nodes.append(child)

    while len(nodes) > 0:
        process_node(nodes.popleft())

    return colors

def _get_connected_numbers(n: int, bits: int):
    """Returns the list of numbers that are one bit flip away from n."""
    return [n ^ (1 << bit) for bit in range(bits)]

Generated solutions

For N = 2

For N = 4

For N = 8