N-Queens in Python

After reading aphyr - Typing the Technical Interview , I wondered how quickly I could solve a similar problem without the aid of type-level magic.

Here's what I came up with after around 10 minutes.

The idea is to recursively solve the problem by adding queens to a (initially empty) partial solution. This partial solution keeps track of the position of the queen in each row.

To find a position for the queen in row i , we compute available (non-attacked) positions for it by starting with a set of all positions and removing ones attacked by queens in previous rows.

For example, a queen in column c of row 0 attacks columns {c-1, c, c+1} in row 1, columns {c-2, c, c+2} in row 2 and so on.

We can then iterate over all available columns and try place the next queens based on the new partial solution. If this does not yield a solution, we continue with the next available column.

One downside of this approach is that large values of n will exhaust the recursion limit.

def search(n, partial, i):
    if i == n:
        return partial

    available = set(range(0, n))

    for j in range(0, i):
        delta = i - j
        c = partial[j]
        available -= {c - delta, c, c + delta}

    for a in available:
        partial[i] = a
        res = search(n, partial, i + 1)

        if res is not None:
            return res

def prettyprint(solution):
    n = len(solution)
    for r in range(0, n):
        row = ['_'] * n
        row[solution[r]] = 'Q'
        print(''.join(row))

if __name__ == '__main__':
    n = 8
    partial = [None] * n
    solution = search(n, partial, 0)
    print(solution)
    print()
    prettyprint(solution)

[0, 4, 7, 5, 2, 6, 1, 3]

Q_______
____Q___
_______Q
_____Q__
__Q_____
______Q_
_Q______
___Q____