nqueens_row.py

class NQueensRow(SearchProblem):
    def __init__(self, N):
        super(NQueensRow, self).__init__()
        self.N             = N
        self.initial_state = ()
        self._actions      = range(N)
        shuffle(self._actions) # randomize actions

    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) # randomize actions at each step
        return tmp

    def result(self, s, a):
        '''Result of applying an action to a state.'''
        b = list(s)     # make mutable to add
        b.append(a)
        return tuple(b) # needs to be immutable again to search

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

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

    def _is_valid(self, s):
        '''Check if a state is valid.'''
        # valid states: any arrangement of N queens with none attacking each other
        # check horizontal, vertical, right/left diagonals
        attacked = self._attacked
        return attacked(s, (1, 1)) and attacked(s, (1,-1)) and \
               attacked(s, (0, 1)) and attacked(s, (1, 0))