NQueens.py

--D. Thiebaut 16:04, 24 January 2010 (UTC)

#! /opt/local/bin/python
# D. Thiebaut
# Usage:
#   nqueens.py N
# where N is the number of queens wanted, i.e. the dimension
# of the board.
#

import sys
import time
QUEEN = -10
EMPTY = 0

class timePerfClass:
    """A class to keep track of processor and user time.
    use by first calling start(), then stop(), then 
    printElapsed()"""

    def __init__(self):
        self.clockStart = 0
        self.clockEnd   = 0
        self.timeStart  = 0
        self.timeEnd    = 0

    def start( self ):
        self.clockStart = time.clock()
        self.timeStart  = time.time()

    def stop( self ):
        self.clockEnd   = time.clock()
        self.timeEnd    = time.time()

    def printElapsed( self ):
        print "processor time:   ", (self.clockEnd-self.clockStart), "sec"
        print "user time     :   ", (self.timeEnd-self.timeStart), "sec"


def makeBoard( N ):
    """create a 2-D array of ints with dimension N
    Returns the 2D array"""
    board = []
    for i in range( N ):
        board.append( [] )
        for j in range( N ):
            board[i].append( EMPTY )
    return board
    
def goHome():
    sys.stderr.write( "\x1b[0;0H" )

def displaySolution( board ):
  list = []
  for i in range( len( board ) ):
      for j in range( len( board ) ):
          if board[ i ][ j ]==QUEEN:
              list.append( str( j ) )
  print "Solution = ", ','.join( list )

def displayBoard( board, home=False ):
    """display the 2D array, showing empty cells as .
    and queens as Q"""
    if home: goHome()
    for i in range( len( board ) ):
        for j in range( len( board ) ):
            if board[i][j]==QUEEN: 
               print 'Q',
            else: 
               print '.',
        print
    displaySolution( board )
    
def markLowerBoard( board, row, col, offset ):
    """Once a queen is positioned at location (row,col), 
    all the cells on a lower diagonal and lower vertical
    of this queen must be marqued as unavailable so that
    another queen cannot be put in this place"""

    N = len( board )
    diagleft = col-1
    diagright = col+1
    # mark all lower rows on diagonals and vertical
    for r in range( row+1, N ):
        if diagleft >=0:  board[r][diagleft] += offset
        if diagright <N:  board[r][diagright] += offset
        board[r][col] += offset
        diagleft  -= 1
        diagright += 1

def tryRow( board, row,  N, display=False ):
    """ put a queen on give row, and recursively try all queens
    on successive rows"""
    if row >= N:
       return True #we found a solution!

    if display:
        displayBoard( board, True )
    for col in range( N ):
        if board[row][col] == EMPTY:
            # put a queen here
            board[ row ][ col ] = QUEEN
            markLowerBoard( board, row, col, +1 )
            ok = tryRow( board, row+1, N )
            if not ok:
                # backtrack
                board[ row ][ col ] = EMPTY
                markLowerBoard( board, row, col, -1 )
            else:
                return True
    return False           
           
       
def main():
    if len( sys.argv ) < 2:
        print "Syntax: nqueens.py N"
        print "        where N is the # of queens"
        return

    #--- start measurement ---
    perf = timePerfClass()
    perf.start()

    #--- get dimension, create board, and solve! ---
    N = int( sys.argv[1] )
    board = makeBoard( N )
    ok = tryRow( board, 0, N )

    #--- stop measurement ---
    perf.stop()
    perf.printElapsed()
    
    if ok:
         print "Success!"
         displayBoard( board )

if __name__=="__main__":
    main()