n-queens-solution-via-gauss.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Implement a console script that solves the standard Queens puzzle:

 * Take an argument – size of the board (N).
 * Place N chess queens on an N×N chessboard so that no two queens threaten each other.
 * Print the number of possible solutions.
 * Print the time that was required to solve the problem.

Bonus 1
 * Take an optional argument that displays an example of a solution or all solutions.

Bonus 2

 * Take an optional argument that will mutate all queens to
   “maharajas” – pieces that can attack as a queen and as a knight
   simultaneously.

Known values: (1, None, None, 2, 10, 4, 40, 92, 352, 724, 2680, 14200)
"""

import argparse
import itertools
import operator
import random
import sys
import time


def solution_sequence(n):
    """This generates all states for given dimension value, searches for
    problem's solutions and return them."""
    fwd = tuple(range(n))
    rev = tuple(reversed(fwd))
    solutions = {}

    # this is a bit faster but not so informative
    # return [state for state in itertools.permutations(range(n), n)
    #         if len(set(map(operator.add, fwd, state))) == len(set(map(operator.add, rev, state))) == n]

    for state in itertools.permutations(range(n), n):
        res1 = set(map(operator.add, fwd, state))
        res2 = set(map(operator.add, rev, state))

        if len(res1) == len(res2) == n:
            solutions[state] = 1

    return solutions.keys()


def show_board(state):
    """Prints a board with given solution on."""
    max_col = len(state)
    for row, value in enumerate(state):
        line = ['Q' if col == value else '+' for col in xrange(max_col)]
        print('\t' + ' '.join(line))


def show_solution(index, state):
    """Prints textual representation of given solution."""
    print('Solution[{}] is {}'.format(index + 1, state))


def arg_parser():
    """Defines the parser's rules."""
    parser = argparse.ArgumentParser(description='N queens problem.')
    parser.add_argument('dimension', type=int, help='dimension of a board')
    parser.add_argument('--debug', action='store_true', help='enabledebug output ')
    parser.add_argument('--board', action='store_true', help='show solution on board ')
    parser.add_argument('--show', choices=['first', 'random', 'all'], help='how to show solutions')
    return parser.parse_args()


if __name__ == '__main__':
    args = arg_parser()

    dimension = args.dimension
    if 0 >= dimension or 1 < dimension < 4:
        print('Error: Invalid input data')
        print('Valid dimensions are 1, 4, 5, 6, 7, 8 and so on')
        sys.exit(1)

    start = time.time()
    solutions = solution_sequence(dimension)
    spent = time.time() - start

    print('{:>10}: {:>8d}'.format('Solutions', len(solutions)))
    print('{:>10}: {:>8.02f} sec'.format('Time', spent))

    if 'all' == args.show:
        for index, state in enumerate(solutions):
            show_solution(index, state)
        if args.board:
            print('Board will show with show=first/random.')
    elif args.show in ('random', 'first'):
        state = random.choice(solutions) if 'random' == args.show else solutions[0]

        print('')
        show_solution(solutions.index(state), state)

        if args.board and all != args.show:
            show_board(state)

sys.exit(0)