nqueens.py

#!/usr/bin/env python3

import numpy as np

def atk(board, x, y):
    h, w = board.shape
    rest = w - x
    atk = np.vstack([np.zeros((y, rest), dtype=int), np.ones(rest, dtype=int), np.zeros((h - 1 - y, rest), dtype=int)])
    atk |= np.eye(N=h, M=rest, k=-y, dtype=int)
    atk |= np.flip(np.eye(N=h, M=rest, k=y-h+1, dtype=int), axis=0)
    atk = np.hstack([np.zeros((h, x), dtype=int), atk])
    return board | atk

# Go column by column, from left to right.
# Return pairs (board, indexes) such as `board`
# contains the temporary state required for solving the problem,
# and `indexes` contains the row indexes (Y coordinates) of the
# queen on the X-th column.
def solve(board, x):
    h, w = board.shape

    if x < w: # we haven't reached the end of the table
        safe = np.argwhere(board[:, x] == 0).ravel()

        for y in safe:
            tmp = atk(board, x, y)
            
            for board2, y2 in solve(tmp, x + 1):
                yield board2, np.hstack([y, y2]).astype(int)

    else:
        yield board, []

def solve_pretty(n):
    board = np.zeros((n, n), dtype=int)

    for _, ix in solve(board, 0):
        yield '\n'.join('.' * i + 'Q' + '.' * (n - i - 1) for i in ix)

for i, solution in enumerate(solve_pretty(8)):
    print('Solution', i + 1)
    print(solution)
    print()