"""
 queens.py

 Working on the N-queens problem.

 A representing boards with one queen per row,
 with queens[row] giving the columns containting a queen.

 >>> queens = [0, 3, 3, 2]
 >>> print board(queens)
    Q . . .
    . . . Q
    . . . Q
    . . Q .
 >>> collisions(queens)   # one vertical pair, one diagonal pair
 2

 Jim Mahoney | April 2013 | MIT License
"""
from random import randint

def make_random_queens(n):
    """ Return random configuruation, an N-element list of queen columns """
    return [randint(0, n) for i in range(n)]

def board(queens):
    """ Return string representation of chess board with queens on it """
    result = ''
    N = len(queens)
    for column in queens:
        row = ['.'] * N
        row[column] = 'Q'
        result += '   ' + ' '.join(row) + '\n'
    return result.rstrip()

def collisions(queens):
    """ Return number of pairs of queens vertically or diagonally apart. """
    # Looking at queens in two different rows, e.g.
    #
    #   . . . . . 
    #   . . Q . .  rowA=1, columnA=2
    #   . . . . .
    #   . . . . Q  rowB=3, columnB=4
    #
    # the two will be vertically aligned if columnA-columnB==0,
    # and will be diagonally aligned if columnA-columnB = +-(rowA-rowB)
    count = 0
    N = len(queens)
    for rowA in range(N-1):                 # rowA : 0 down to next-to-last
        for rowB in range(rowA + 1, N):     # rowB : belowA to last
            row_diff = rowB - rowA
            col_diff = queens[rowB] - queens[rowA]
            if col_diff in (0, row_diff, -row_diff):
                count += 1
    return count

if __name__ == "__main__":
    import doctest
    doctest.testmod()