csp_queen.py

from random import choice
from time import time

N = int(input("N > "))


class COLORS:
    GREEN = '\033[92m'
    RED = '\033[91m'
    ENDC = '\033[0m'


class Board:
    def __init__(self):
        """
            initialize the board
            _queen_rows: a list of sets, each set contains the positions of the queens in the same row
            _queen_posdiag: a list of sets, each set contains the positions of the queens in the same positive diagonal
            _queen_negdiag: a list of sets, each set contains the positions of the queens in the same negative diagonal
        """
        self._queen_rows = {i: set() for i in range(1, N+1)}
        self._queen_posdiag = {i: set() for i in range(1, 2 * N)}
        self._queen_negdiag = {i: set() for i in range(1, 2 * N)}

    def set_queen(self, x, y, constraints):
        """
            set the queen at (x, y) with the constraints
            update the constraints with calculating the conflicts
        """
        combined = self._queen_rows[y] | \
            self._queen_posdiag[y + (x-1)] | \
            self._queen_negdiag[y + (N - x)]

        for i in combined:
            constraints[i] += 1

        self._queen_rows[y].add(x)
        self._queen_posdiag[y+(x-1)].add(x)
        self._queen_negdiag[y+(N - x)].add(x)

        constraints[x] = len(combined)

    def remove_queen(self, x, y, constraints):
        """
            remove the queen at (x, y)
            update the constraints with calculating the conflicts
        """
        combined = self._queen_rows[y] | \
            self._queen_posdiag[y + (x-1)] | \
            self._queen_negdiag[y + (N - x)]

        for i in combined:
            constraints[i] -= 1

        self._queen_rows[y].remove(x)
        self._queen_posdiag[y+(x-1)].remove(x)
        self._queen_negdiag[y+(N - x)].remove(x)

        constraints[x] = 0
        return

    def get_num_conflicts(self, x, y):
        """
            calculate the number of conflicts for the queen at (x, y)
        """
        combined = self._queen_rows[y] | \
            self._queen_posdiag[y + (x-1)] | \
            self._queen_negdiag[y + (N - x)]
        return len(combined)

    def get_least_conflicts_y(self, x, possible):
        """
            calculate the least number of conflicts for the queen at (x, possible=[y1, ..., yn])
        """
        conflict_list, min_count = [possible[0]
                                    ], self.get_num_conflicts(x, possible[0])
        for i in possible[1:]:
            count = self.get_num_conflicts(x, i)
            if min_count > count:
                min_count = count
                conflict_list = [i]
            elif min_count == count:
                conflict_list.append(i)

        return choice(conflict_list)


class CSP:
    def __init__(self):
        """
            initialize the CSP
        """
        self.board = Board()
        self.variables = None
        self.domains = None
        self.constraints = None

    def print_board(self):
        def create_board(self):
            b = [['-' for i in range(N)] for j in range(N)]
            for key, value in self.domains.items():
                if self.constraints[key] > 0:
                    b[value-1][key - 1] = COLORS.RED + 'Q' + COLORS.ENDC
                else:
                    b[value - 1][key - 1] = COLORS.GREEN + 'Q' + COLORS.ENDC
            return b
        for x in create_board(self):
            for y in x:
                print(y, end=' ')
            print()
        print()

    def backtrack(self, max_step=9000):
        self.variables = [i for i in range(1, N+1)]
        self.constraints = {i: 0 for i in self.variables}

        _var = self.variables.copy()
        disallow_loc = {i: set() for i in self.variables}

        start_solve = time()

        y = choice(_var)
        self.domains = {1: y}
        _var.remove(y)
        self.board.set_queen(1, y, self.constraints)
        x = 2
        step = 1

        while x <= N and step <= max_step:
            variable = set(_var) - disallow_loc[x]
            possibles = list(filter(lambda x: x[1] <= 1, map(
                lambda y: (y, self.board.get_num_conflicts(x, y)), variable)))
            possibles.sort(key=lambda x: x[1])
            for possible in possibles:
                y, _ = possible
                self.board.set_queen(x, y, self.constraints)
                if self.constraints[x] > 0:
                    disallow_loc[x].add(y)
                    self.board.remove_queen(x, y, self.constraints)
                else:
                    self.domains[x] = y
                    _var.remove(y)
                    x += 1
                    break
            else:
                if x > 1:
                    x -= 1
                    for xi in range(x+1, N+1):
                        disallow_loc[xi].clear()
                    disallow_loc[x].add(self.domains[x])
                    self.board.remove_queen(
                        x, self.domains[x], self.constraints)
                    _var.append(self.domains[x])
                    del self.domains[x]
                    if disallow_loc[x] == set(self.variables):
                        print("backtrack problem")
                        exit(0)
            step += 1

        print("step number : {}".format(step))
        self.print_board()
        print("run time {:.4f}s".format(time() - start_solve))
        if self.variables != list(self.domains.keys()):
            print(COLORS.RED + "unsolved" + COLORS.ENDC)
        else:
            print(COLORS.GREEN + "solved" + COLORS.ENDC)


CSP().backtrack()