Generating solutions to the 8 Queens Puzzle

2006-07-21, Comments

Here’s a Python solution to the 8 Queens puzzle. I think that generator functions are perfectly suited to this kind of thing: they are easy to code and easy to use, they yield results as soon as they can, and clients have full control over when to stop them yielding results.

n-queens.py
class Queen:
    """A queen on a chess board.
    """
    def __init__(self, row, col):
        self.row, self.col = row, col

    def __str__(self):
        return "(%d, %d)" % (self.row, self.col)

    def safe_from(self, queens):
        """Is this queen safe from the input list of queens?
        """
        return not any(self.attacked_by(queen) for queen in queens)

    def attacked_by(self, queen):
        """Is this queen attacked by the input queen?
        """
        r, c = self.row, self.col
        qr, qc = queen.row, queen.col
        def same_row():
            return r == qr
        def same_column():
            return c == qc
        def same_diagonal():
            return r + c == qr + qc or r - c == qr - qc
        return same_row() or same_column() or same_diagonal()

def queens_puzzle(board_size):
    """Generate solutions to the Queens puzzle for the input board size.
    """
    def queen_columns(cols):
        if cols == 0:
            yield list()
        else:
            cols -= 1
            for queens in queen_columns(cols):
                for row in range(board_size):
                    new_queen = Queen(row, cols)
                    if new_queen.safe_from(queens):
                        yield queens + [new_queen]
    for queens in queen_columns(board_size):
        yield queens

def outer_join(s, seq):
    return "%s%s%s" % (s, s.join(seq), s)

def print_solution(board_size, queens):
    line = outer_join("+", "-" * board_size)
    for row in range(board_size):
        print line
        print outer_join("|",
                         [" Q"[q.row == row] for q in queens])
    print line
    print

try:
    print """\
Queens puzzle solver.
Puts N chess queens on an N by N chessboard such that none of
them is able to capture any other.

At the prompt, enter the board size or Q to quit.
"""
    while True:
        board_size = int(raw_input("board size (Q to quit)? "))
        if board_size <= 0:
            print "Please enter a positive board-size."
        else:
            if board_size > 16:
                print ("That's a tough one but I'll try my best. "
                       "Press CTRL-C to interrupt.")
            try:
                solutions = -1
                for solutions, queens in enumerate(queens_puzzle(board_size)):
                    print_solution(board_size, queens)
                solutions += 1 # Users count from 1, not 0
                print "There %s %d solution%s to the %d Queen%s puzzle.\n" % (
                    "is" if solutions == 1 else "are",
                    solutions, "s"[solutions==1:],
                    board_size,"s"[board_size==1:])
            except KeyboardInterrupt:
                print "Execution interrupted."
except ValueError, v:
    print "Bye for now!"