nqueens_square.py

class NQueensSquare(SearchProblem):
    def __init__(self, N):
        super(NQueensSquare, self).__init__()
        self.N = N
        self.initial_state = tuple([tuple([0 for i in range(N)]) for j in range(N)])
        self._actions      = [(i, j) for i in range(N) for j in range(N)]

    def actions(self, s):
        '''Possible actions from a state.'''
        # generate every possible state then filter out invalid
        tmp = [a for a in self._actions if self._is_valid(self.result(s, a))] 
        shuffle(tmp)
        return tmp

    def result(self, s, a):
        '''Result of applying an action to a state.'''
        (i, j) = a
        b      = list(s)    # make immutable object to place queen
        tmp    = list(b[i])    
        tmp[j] = 1          # place queen to square (i,j) in board

        s    = list(s)      # extra steps for immutability
        s[i] = tuple(tmp)   # required by A* search in simpleai
        return tuple(s)

    def is_goal(self, state):
        '''Goal: N queens on board and no queen under attack.'''
        return self._num_queens_on_board(state) == self.N

    def heuristic(self, state):
        return self.N - self._num_queens_on_board(state)