N-queens, initial implementation

nqueens.py

''' Solves the n-queens problem, displaying all solutions as we find them.

Is currently ordered (shows all permutations of queen placements), which is
extremely non-optimal.

Note: this is an original algorithm I worked out on paper fwiw; the core idea
is representing queen placement with 3 1d arrays (not one 2d array), and a 
recursive approach where we execute and then unexecute placement operations
to reuse the same memory. 

The only memory intensive operation is when we paint the board (I/O + string ops).
'''
def nqueens (N):
    # Three arrays:
    #   X:  signals empty (0) / nonempty ([1, n]) rows
    #   Y:  signals empty (0) / nonempty ([1, n]) columns
    #   Q:  list of queen positions (x, y)
    #
    # The following holds for each nth [0, N) queen at (x, y):
    #   X[x] == n+1
    #   Y[y] == n+1
    #   Q[n] == (x, y)
    #
    X, Y, Q = [ 0 ] * N, [ 0 ] * N, [ (-1, -1) ] * N

    ''' Displays the board to stdout. Used to display solutions when we find them. '''
    def paint ():
        board = [ [ 0 ] * N for i in range(N) ]
        for i, (x, y) in enumerate(Q):
            board[y][x] = i+1
        for row in board:
            print(row)
        print("")

    ''' Place the nth queen at (x, y). '''
    def put (n, x, y):
        assert(X[x] == 0)
        assert(Y[y] == 0)
        assert(Q[n] == (-1, -1))
        X[x] = n+1
        Y[y] = n+1
        Q[n] = (x, y)

    ''' Undo a put operation, restoring the board to its previous state. '''
    def unput (n, x, y):
        assert(X[x] == n+1)
        assert(Y[y] == n+1)
        assert(Q[n] == (x, y))
        X[x] = 0
        Y[y] = 0
        Q[n] = (-1, -1)

    ''' Generator that yields all valid (x, y) placement positions, given n placed
    queens and the board's current state and queen placement constraints:
    1)  A queen may not be placed in the same row as another queen (check X[x] == 0)
    2)  A queen may not be placed in the same column as another queen (check Y[y] == 0)
    3)  A queen may not be placed diagonal to another queen:
        voilated if abs(qx - x) == abs(qy - y) for any (qx, qy) in Q[:n]
    '''
    def pick_xy (n):
        def constrain_diagonal (x, y):
            for qx, qy in Q[:n]:
                # if qx == -1:
                #     continue
                if abs(qx - x) == abs(qy - y):
                    return False
            return True

        for x in range(0, N):
            if X[x] != 0:
                continue
            for y in range(0, N):
                if Y[y] != 0:
                    continue
                if constrain_diagonal(x, y):
                    yield x, y

    ''' Recursive function to place the nth queen iff possible, for n in [0, N) '''
    def place_queen (n = 0):
        if n >= N:
            # Finished placing N queens: print results and continue
            paint()
            return

        # Core algorithm: for each valid x / y position:
        #   1) place at x, y
        #   2) recursively explore this placement
        #   3) undo placement and apply other positions.
        #
        # Is fast-ish b/c we can exclude rows / columns quickly with the X / Y 
        # arrays, only need to do up to N diagonal comparisons, and state is
        # cached in generators (pick_xy).
        #
        # Is slow b/c we're using ordered queen placement which turns out to be
        # highly redundant, so we're displaying / exploring a lot of duplicate
        # boards where only the order of the queens changes. I'll try to fix
        # this on the second iteration, either using an improved algorithm or
        # some kind of hashing.
        #
        for x, y in pick_xy(n):
            put(n, x, y)
            place_queen(n+1)
            unput(n, x, y)
    place_queen()

if __name__ == '__main__':
    nqueens(4)

output.txt

[0, 0, 3, 0]
[1, 0, 0, 0]
[0, 0, 0, 4]
[0, 2, 0, 0]

[0, 0, 4, 0]
[1, 0, 0, 0]
[0, 0, 0, 3]
[0, 2, 0, 0]

[0, 0, 2, 0]
[1, 0, 0, 0]
[0, 0, 0, 4]
[0, 3, 0, 0]

[0, 0, 2, 0]
[1, 0, 0, 0]
[0, 0, 0, 3]
[0, 4, 0, 0]

[0, 0, 4, 0]
[1, 0, 0, 0]
[0, 0, 0, 2]
[0, 3, 0, 0]

[0, 0, 3, 0]
[1, 0, 0, 0]
[0, 0, 0, 2]
[0, 4, 0, 0]

[0, 2, 0, 0]
[0, 0, 0, 4]
[1, 0, 0, 0]
[0, 0, 3, 0]

[0, 2, 0, 0]
[0, 0, 0, 3]
[1, 0, 0, 0]
[0, 0, 4, 0]

[0, 3, 0, 0]
[0, 0, 0, 4]
[1, 0, 0, 0]
[0, 0, 2, 0]

[0, 4, 0, 0]
[0, 0, 0, 3]
[1, 0, 0, 0]
[0, 0, 2, 0]

[0, 3, 0, 0]
[0, 0, 0, 2]
[1, 0, 0, 0]
[0, 0, 4, 0]

[0, 4, 0, 0]
[0, 0, 0, 2]
[1, 0, 0, 0]
[0, 0, 3, 0]

[0, 1, 0, 0]
[0, 0, 0, 4]
[2, 0, 0, 0]
[0, 0, 3, 0]

[0, 1, 0, 0]
[0, 0, 0, 3]
[2, 0, 0, 0]
[0, 0, 4, 0]

[0, 1, 0, 0]
[0, 0, 0, 4]
[3, 0, 0, 0]
[0, 0, 2, 0]

[0, 1, 0, 0]
[0, 0, 0, 3]
[4, 0, 0, 0]
[0, 0, 2, 0]

[0, 1, 0, 0]
[0, 0, 0, 2]
[3, 0, 0, 0]
[0, 0, 4, 0]

[0, 1, 0, 0]
[0, 0, 0, 2]
[4, 0, 0, 0]
[0, 0, 3, 0]

[0, 0, 3, 0]
[2, 0, 0, 0]
[0, 0, 0, 4]
[0, 1, 0, 0]

[0, 0, 4, 0]
[2, 0, 0, 0]
[0, 0, 0, 3]
[0, 1, 0, 0]

[0, 0, 2, 0]
[3, 0, 0, 0]
[0, 0, 0, 4]
[0, 1, 0, 0]

[0, 0, 2, 0]
[4, 0, 0, 0]
[0, 0, 0, 3]
[0, 1, 0, 0]

[0, 0, 4, 0]
[3, 0, 0, 0]
[0, 0, 0, 2]
[0, 1, 0, 0]

[0, 0, 3, 0]
[4, 0, 0, 0]
[0, 0, 0, 2]
[0, 1, 0, 0]

[0, 0, 1, 0]
[2, 0, 0, 0]
[0, 0, 0, 4]
[0, 3, 0, 0]

[0, 0, 1, 0]
[2, 0, 0, 0]
[0, 0, 0, 3]
[0, 4, 0, 0]

[0, 0, 1, 0]
[3, 0, 0, 0]
[0, 0, 0, 4]
[0, 2, 0, 0]

[0, 0, 1, 0]
[4, 0, 0, 0]
[0, 0, 0, 3]
[0, 2, 0, 0]

[0, 0, 1, 0]
[3, 0, 0, 0]
[0, 0, 0, 2]
[0, 4, 0, 0]

[0, 0, 1, 0]
[4, 0, 0, 0]
[0, 0, 0, 2]
[0, 3, 0, 0]

[0, 3, 0, 0]
[0, 0, 0, 4]
[2, 0, 0, 0]
[0, 0, 1, 0]

[0, 4, 0, 0]
[0, 0, 0, 3]
[2, 0, 0, 0]
[0, 0, 1, 0]

[0, 2, 0, 0]
[0, 0, 0, 4]
[3, 0, 0, 0]
[0, 0, 1, 0]

[0, 2, 0, 0]
[0, 0, 0, 3]
[4, 0, 0, 0]
[0, 0, 1, 0]

[0, 4, 0, 0]
[0, 0, 0, 2]
[3, 0, 0, 0]
[0, 0, 1, 0]

[0, 3, 0, 0]
[0, 0, 0, 2]
[4, 0, 0, 0]
[0, 0, 1, 0]

[0, 3, 0, 0]
[0, 0, 0, 1]
[2, 0, 0, 0]
[0, 0, 4, 0]

[0, 4, 0, 0]
[0, 0, 0, 1]
[2, 0, 0, 0]
[0, 0, 3, 0]

[0, 2, 0, 0]
[0, 0, 0, 1]
[3, 0, 0, 0]
[0, 0, 4, 0]

[0, 2, 0, 0]
[0, 0, 0, 1]
[4, 0, 0, 0]
[0, 0, 3, 0]

[0, 4, 0, 0]
[0, 0, 0, 1]
[3, 0, 0, 0]
[0, 0, 2, 0]

[0, 3, 0, 0]
[0, 0, 0, 1]
[4, 0, 0, 0]
[0, 0, 2, 0]

[0, 0, 4, 0]
[2, 0, 0, 0]
[0, 0, 0, 1]
[0, 3, 0, 0]

[0, 0, 3, 0]
[2, 0, 0, 0]
[0, 0, 0, 1]
[0, 4, 0, 0]

[0, 0, 4, 0]
[3, 0, 0, 0]
[0, 0, 0, 1]
[0, 2, 0, 0]

[0, 0, 3, 0]
[4, 0, 0, 0]
[0, 0, 0, 1]
[0, 2, 0, 0]

[0, 0, 2, 0]
[3, 0, 0, 0]
[0, 0, 0, 1]
[0, 4, 0, 0]

[0, 0, 2, 0]
[4, 0, 0, 0]
[0, 0, 0, 1]
[0, 3, 0, 0]

[Finished in 0.4s]