N-Queens Solver

queens.py

#!/usr/bin/env python
import random

def queens(n):
    """Places n queens on an n*n chessboard so that each is safe from all others.

    This solution uses a depth-first backtracking algorithm.

    Argument
    ========
    n -- the number of queens

    Algorithm
    =========
    A depth-first backtracking algorithm traverses a tree,
    going as 'deep' as possible into the tree until there is no way
    forward. At that point, it backtracks to its previous point and tries
    an alternative option.

    In this program, each level of the tree corresponds with the placement
    of one queen, so the algorithm is successful when it reaches a leaf at
    depth n.

    The entire tree is not stored in memory; after each iteration, any given
    space is only found in one place. The only exceptions to this are spaces
    with queens on them, which are both in the 'board' and in 'current_spaces'.

    The active branch of the tree (the list of points with queens on them)
    is stored in 'current_spaces'. Potential children of the current leaf
    are found in 'available'. Spaces that are unavailable because they are
    occupied or in an attack lane are stored in the 'board'. The algorithm
    avoids traversing permutations of unsuccessful branches by storing the
    roots of these branches in a separate dict, 'already_tried', and excluding
    them from selection. These already-tried spaces are not on the board,
    even if they are in an attack lane.

    Each of these structures is updated any time a queen is added to or
    removed from the board. Additionally, the first time a queen is placed
    on a space, that space's attack lanes are generated and saved in
    'attack_lanes' for future use.

    Any solution to n-queens places one queen per row. To reduce the number
    of spaces to try with each leaf, the algorithm places queens one row at
    a time, beginning with the first, then the second, etc.


    Data Structures
    ===============

    * attack_lanes (a dict)
    * structure:    { space : [attack_lanes]}
                    {(0, 0) : [ (0, 1), (0, 2) ... ] }
        Contains attack lanes for each space. A queen may attack any space
        in its row or column or along either of its two diagonals. Lists
        of such points are generated as needed and stored in the attack_lanes
        dictionary for re-use.

    * board (a dict)
    * structure:    { space : queen-number, space : None }
                    { (0, 0) : 4, (0, 1) : 4, (1, 7) : None }
        Contains all spaces on the board, except spaces already tried and
        rejected for use with this branch. Spaces containing a queen or in
        one of her attack lanes are marked with the number of that queen.
        Other spaces are denoted by None.

    * current_spaces (a list)
    * structure:    [ space ]
                    [ (0,0), (1, 7) ... ]
        Contains all spaces currently occupied by a queen.

    * already_tried (a dict)
    * structure:    { space : parent-queen-number }
                    { (6, 2) : 6 }
        Contains spaces that were already tried and are not compatible with
        at least one queen on the board. These are the immediate children
        of any members of current_spaces.
    """

    # check input
    if n < 1 or n in (2,3):
        raise ValueError('The n-queens problem is not solvable '
                         'for n < 1, n=2, or n=3')

    # initialize data structures
    current_spaces = []
    attack_lanes = {}
    already_tried = {}
    board = {}
    for i in range(0,n):
        for j in range(0,n):
            board[(i,j)] = None

    def make_attack_lanes(space):
        """function for generating attack lanes"""
        row = space[0]
        col = space[1]
        attackable = []

        # add row & col attack lanes
        for i in range(0,n):
            attackable.append((i,col))
        for j in range(0,n):
            attackable.append((row,j))

        # add diagonal attack lanes
        for i in range(0,n):
            j1 = i - row + col
            j2 = -(i - row) + col

            for j in (j1, j2):
                if 0 <= j < n:
                    attackable.append((i, j))
        return list(set(attackable))

    # place the queens!
    q = 0
    run = 0
    while q < n:
        # get a list of open spaces in the current row
        available = {k : v for k,v in board.iteritems()
                     if v == None and k[0] == q}

        # space is available - add queen to board
        if available:
            space = random.choice(available.keys())
            board[space] = q

            # add attack lanes to board
            try:
                als = attack_lanes[space]
            except KeyError:
                attack_lanes[space] = als = make_attack_lanes(space)
            for s in als:
                try:
                    if board[s] is None:
                        board[s] = q
                except KeyError:
                    # the space is in already_tried
                    pass

            current_spaces.append(space)
            q += 1
            run += 1

        # no space is available - need to backtrack
        else:
            q -= 1

            # move queen from board to already_tried
            if current_spaces:
                bad_space = current_spaces.pop(q)
                del board[bad_space]
                already_tried[bad_space] = q-1

            # remove queen's attack lanes from board
            for k,v in board.items():
                if v == q:
                    board[k] = None

            # move any children in already_tried back to board
            if already_tried:
                for k,v in already_tried.items():
                    if v == q:
                        board[k] = None
                        del already_tried[k]
            run += 1

    # display output - queens are 1's, others are 0's
    display_board = [[0 for x in range(n)] for y in range(n)]
    for space in current_spaces:
        row = space[0]
        col = space[1]
        display_board[row][col] = 1

    print "The board is: "
    for row in display_board:
        print row
    print ("This run took " + str(run) +
           " steps to place " + str(n) + " queens.")
    return

n = int(raw_input("Enter an integer greater than 3: "))
queens(n)