nqueens_swap.py

class NQueensSwap(SearchProblem):
    def __init__(self, N):
        self.N = N
        init               = [i for i in range(N)]
        shuffle(init)   # randomize the initial array
        self.initial_state = tuple(init)    # states must me immutable
        self._actions      = [ (i, j) for i in range(N) 
                                      for j in range(i,N) 
                                          if i != j ]

    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 board mutable to swap queens
        (i, j)     = a
        b[i], b[j] = b[j], b[i] # swap queens
        return tuple(b)         # make immutable again to search

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

    def heuristic(self, state):
        # scan the state horizontally, vertically and diagonally (left & right)
        # return the number of attacks summed for all queens
        nattacks = self._num_attacks
        return nattacks(state, (0, 1)) + nattacks(state, (1, 0)) + \
               nattacks(state, (1,-1)) + nattacks(state, (1, 1))

    def _is_valid(self, s):
        '''Check if a state is valid.'''
        # all states are valid by default
        return True