538-riddler/2020_08_01/main.py

import numpy as np


def fill_and_mark(t, stud, pos):
    t[:, pos] = 'x'
    t[stud, :] = 'x'
    ys, xs = np.diag_indices_from(t)
    xs = (xs + pos) % len(t)
    ys = (ys + stud) % len(t)
    t[ys, xs] = 'x'
    t[stud, pos] = 'o'


def main():
    n = 9

    t_init = np.full((n, n), ' ', dtype=str)
    fill_and_mark(t_init, 0, 0)

    sols = []
    stack = [(t_init, 1)]
    while stack:
        t, curr_stud = stack.pop()

        if curr_stud == n:
            sols.append(t)
            continue

        for pos in np.argwhere(t[curr_stud] == ' '):
            t_new = t.copy()
            fill_and_mark(t_new, curr_stud, pos)
            stack.append((t_new, curr_stud+1))

    print('number of solutions: ', len(sols))
    for solution in sols:
        print(solution)


if __name__ == "__main__":
    main()