Solve and print all 8 (or more) queens problems using backtracking

8_queens.py

# Author: cybojanek
# Board size
SIZE=8
# ASCI value for queen
QUEEN='Q'


def number_to_row_permutation(n):
    """Retun the nth permutation of a row

    Arguments:
    n - if -1, then return all zeros
    else if [0,SIZE-1], set that field to QUEEN

    Return:
    array of length SIZE, with values of ' ' or QUEEN

    """
    row = [' ' for x in range(SIZE)]
    if n == -1:
        return row
    row[n] = QUEEN
    return row


def print_board(board):
    """Print the board to stdout

    Arguments:
    board - SIZE length array of column numbers

    """
    BLACK_TEMP = '\033[37;40m%-2s\033[0m'
    WHITE_TEMP = '\033[30;47m%-2s\033[0m'
    board = [number_to_row_permutation(x) for x in board]
    for i, row in enumerate(board):
        print ''.join([BLACK_TEMP % col if (i + j) % 2 == 0 else WHITE_TEMP % col for j, col in enumerate(row)])


def construct_candidates(board, row):
    """Given a board, construct all the possible candidates for the next row

    Runtime: O(row)

    Arguments:
    board - SIZE length array of column numbers
    row - row number for which we're returning candidates

    Return:
    Array of possible column positions

    """
    # First row can do anything
    if row == 0:
        return [x for x in range(SIZE)]
    # Otherwise, we need to remove initial conflicts
    all_choices = set([x for x in range(SIZE)])
    taken = set()
    # Remove any column conflicts
    for x in range(0, row):
        # Remove from column
        q_index = board[x]
        taken.add(q_index)
        # Right diagonal
        diag_right = q_index + (row - x)
        if diag_right < SIZE:
            taken.add(diag_right)
        # Left diagonal
        diag_left = q_index - (row - x)
        if diag_left >= 0:
            taken.add(diag_left)
    return [x for x in all_choices.difference(taken)]


solutions = 0
def backtrack(board, row):
    # Keep track of number of solutions
    global solutions
    # Reaching the last one, means we found a viable combination
    if row == SIZE - 1:
        solutions += 1
        print "#" * (2 * SIZE), solutions
        print_board(board)
    else:
        # Target next row
        row = row + 1
        for candidate in construct_candidates(board, row):
            board[row] = candidate
            # To step through, uncomment
            #print_board(board)
            #raw_input()
            backtrack(board, row)
        # We're going back, so reset our row
        board[row] = -1


board = [-1 for x in range(SIZE)]
backtrack(board, -1)