ВИ Тест 4: N-Queens Problem (CSP)

from constraint import *

if __name__ == '__main__':
    problem = Problem(BacktrackingSolver())

    num_queens = int(input())

    queens = range(1, num_queens + 1)
    domain = [(i, j) for i in range(num_queens) for j in range(num_queens)]

    problem.addVariables(queens, domain)

    # ----------------Constraints begin here----------------
    for queen1 in queens:
        for queen2 in queens:
            if queen1 < queen2:
                # Diagonal constraints
                problem.addConstraint(lambda q1, q2: (abs(q1[0] - q2[0]) != abs(q1[1] - q2[1])), (queen1, queen2))
                # Row and column constraint
                problem.addConstraint(lambda q1, q2: q1[0] != q2[0] and q1[1] != q2[1], (queen1, queen2))
    # ------------------------------------------------------
    # If there are less than 7 queens, print only the number of solutions
    if num_queens < 7:
        print(len(problem.getSolutions()))
    else:
        print(problem.getSolution())