n_queens.py

#!/usr/bin/env python                                                                           

from __future__ import division, print_function
import numpy as np

np.random.seed(0)

n = 8
T = 0.01
maxsteps = 100000

# -- Algorithm begins here --

# Queen object used to detect collisions (two queens occupying the same
# square) and clashes (two queens that could take one another).
class Queen(object):
    boardsize = n

    def __init__(self, pos):
        self.pos = pos
    
    # Helper method to determine whether move is on the board
    def is_valid(self, pos):
        if np.any(pos < 0) or np.any(pos >= self.boardsize):
            return False
        return True

    # Make pos a property so we can recalculate the diagonal values
    # on the fly.
    # One queen can take another if they both have the same value for
    # x + y (posdiag) or x - y (negdiag), so we calculate this.
    def set_pos(self, pos):
        if not self.is_valid(pos):
            raise ValueError("Invalid position for Queen {}!".format(pos))
        self._pos = np.array(pos, dtype=int)
        self.posdiag = np.sum(pos)
        self.negdiag = pos[0] - pos[1]

    def get_pos(self):
        return self._pos

    pos = property(get_pos, set_pos)

    def __eq__(self, other):
        if np.all(self.pos == other.pos):
            return True
        return False

    # Do two queens clash?
    def clash(self, other):
        if (self.pos[0] == other.pos[0] or
                self.pos[1] == other.pos[1] or
                self.posdiag == other.posdiag or
                self.negdiag == other.negdiag):
            return True
        return False

    # Returns a list of queens at every possible position this
    # queen can move to
    def get_valid_moves(self):
        moves = []
        for i in range(self.boardsize):
            moves.append(np.array([i, self.pos[1]], dtype=int))
            moves.append(np.array([self.pos[0], i], dtype=int))
            moves.append(np.array([i, self.posdiag - i], dtype=int))
            moves.append(np.array([i, i - self.negdiag], dtype=int))

        valid_moves = []
        for move in moves:
            if not self.is_valid(move):
                continue
            queen = Queen(move)
            if queen != self:
                valid_moves.append(queen)

        return valid_moves

    def __repr__(self):
        return self.pos.__repr__()

# The "energy" of the system is an integer which represents the number of
# clashes on the board. We only count each clash once.
E = 0

queens = []

# Initialize the grid randomly & calculate energy
for i in range(n):
    while True:
        redo = False
        queen = Queen(np.random.randint(n, size=2))
        for other in queens:
            # Two queens occupying the same square is FORBIDDEN
            if queen == other:
                redo = True
                break
            if queen.clash(other):
                E += 1
        if redo:
            continue
        queens.append(queen)
        break

# Print info about our initial board configuration
print(queens)
print(E)
grid = np.zeros((n, n), dtype=bool)
for queen in queens:
    grid[tuple(queen.pos)] = True
print(np.array(grid, dtype=int))


# begin the moves
for i in range(maxsteps):
    accept = True
    dE = 0
    # Pick a random queen to move
    nq = np.random.randint(n)
    queen = queens[nq]
    # Pick a random move for that queen
    move = np.random.choice(queen.get_valid_moves())
    for other in queens:
        # If the move would land on top of another queen, reject it
        if move == other:
            accept = False
            break
        # Calculate the change in energy by counting the clashes between
        # the old position and all other queens, and then between the new
        # position and all other queens
        if queen == other:
            continue
        if queen.clash(other):
            dE -= 1
        if move.clash(other):
            dE += 1
    
    if not accept:
        continue
    
    # If the new move does not land on top of another queen, we accept
    # the move with a probability given by the Metropolis-Hastings algorithm
    prob = min(1, np.exp(-dE/T))
    if np.random.random() > prob:
        accept = False
    
    # If we reject the move, then simply move on
    if not accept:
        continue

    # Update the energy and the positions of the queens
    E += dE
    queens[nq] = move

    # If the energy is 0, then there are no clashes and we've found a solution
    if E == 0:
        print("Found a solution!")
        break

# print some info about the solution
print(i)
print(queens)
print(E)
grid = np.zeros((n, n), dtype=bool)
for queen in queens:
    grid[tuple(queen.pos)] = True
print(np.array(grid, dtype=int))