n_queens_problem.py

"""Solves the 8-queens problem (and its N-queens generalization.)"""

# It is clear that since no two queens can occupy the same column, and there are
# as many columns as queens to place, then each column must contain one and only one
# queen.  The only question is, on which row to place each queen?

# A "left-board" is a way to represent a partial solution to the n-queens problem.
# It is simply a list of integers, each indicating the row occupied by the queen
# in each column. Because it is a partial solution, it can have fewer than 8 elements,
# or even be empty. For example, [0, 2] indicates a that the queen in the left-most
# column is at the top of the board, in the zero-th row, and the queen in the second
# column is two rows, down, etc. The variable qs, short for "queens" is often used
# to hold a left-board.

class QueensError(Exception):
    """custom exception used to signal the end of recursion with no solution found."""
    pass

def is_free(qs, y):
    """given a left-board with no violations, determine if the row y in the next column is free."""
    
    x = len(qs) 
    for qx, qy in enumerate(qs):
        # horizontal
        if qy == y:
            return False
        # diagonal
        dx = x - qx
        dy = y - qy
        if dx == dy or dx == -dy:
            return False
    return True

def print_board(qs, n):
    """output a left-board of queens on an NxN grid."""
    
    for x in range(n):
        for y in range(n):
            print 'Q' if x < len(qs) and qs[x] == y else ' ',
        print

def nqs(qs, n, s):
    """given a left-board of queens, on a board of NxN, anding considering moves only above the s-th row,
    return a complete board of queens or throw a QueensError."""
    
    for y in range(s, n):
        if is_free(qs, y):
            qs2 = qs + [y]
            if len(qs2) < n:
                return nqs(qs2, n, 0)
            else:
                return qs2
    if len(qs) == 0:
        raise QueensError("unable to complete board.")
    last = qs[-1]
    qs2 = qs[:-1]
    # TODO flatten this into a loop instead of using tail-recursion to avoid Python's
    # very low recursion limit.
    return nqs(qs2, n, last+1)


def n_queens(n=8):
    """solve the "N-queens" generalization of the eight queens problem and print the solution."""
    
    try:
        qs = nqs([], n, 0)
        print_board(qs, n)
    except QueensError:
        print 'no solution.'
    except RuntimeError:
        print 'too much recursion'


if __name__ == '__main__':
    # command line version: print the solution to stdout for the given
    # size, defaulting to 8 if no size is explicitly given
    n = 8

    # use a different sized board than 8x8 if the first arg is an integer
    import sys
    if len(sys.argv) > 1:
        try:
            n = int(sys.argv[1])
        except:
            pass

    n_queens(n)