Python Arena
Step 4 of 6

The N queens

A chess queen can move any number of squares along its row, its column or either of its two diagonals, and it attacks every square it could move to. The N queens puzzle asks you to place n queens on an n × n board so that none of them attacks another. On a real 8 × 8 chessboard there are 92 ways to do it, and the puzzle is a favorite of interviewers who want to see backtracking.

Making whole boards and then testing them wastes almost all its time. Eight queens can stand on 64 squares in about 4.4 billion ways. No two queens can share a row, so give each row exactly one: that still leaves 8⁸ = 16,777,216 boards, a column for each of the 8 queens. Ask for different columns too, and you're down to the 8! = 40,320 orderings that itertools.permutations(range(8)) makes, but n! climbs quickly as well: over 3.6 million boards for n = 10 and almost 40 million for n = 11. Backtracking does better. It fills the board from the top row down, and a queen goes on a square only if no queen placed so far attacks it. A row where every square is attacked means the rows above it can't be part of a solution, so the search goes back up a row and moves that queen along. For n = 8 the whole search takes 2,057 calls.

Diagonals in O(1)

To check a square quickly, keep three sets: the columns that have a queen, and the two kinds of diagonals. Along a diagonal that goes down to the right, row - col stays the same. Along one that goes down to the left, row + col does:

for row in range(4):
    diag_keys = "".join(f"{row - col:3}" for col in range(4))
    anti_keys = "".join(f"{row + col:3}" for col in range(4))
    print(diag_keys, "   ", anti_keys)
Output
  0 -1 -2 -3       0  1  2  3
  1  0 -1 -2       1  2  3  4
  2  1  0 -1       2  3  4  5
  3  2  1  0       3  4  5  6

So the square (row, col) is attacked when col in cols, row - col in diag or row + col in anti: three O(1) lookups in sets, instead of a loop over every queen placed so far.

The search

def place(row, n, queens, cols, diag, anti):
    if row == n:                       # a queen in every row: a solution
        for col in queens:
            print("." * col + "Q" + "." * (n - col - 1))
        print()
        return 1
    found = 0
    for col in range(n):
        if col in cols or row - col in diag or row + col in anti:
            continue                   # attacked: prune this branch
        queens.append(col)             # choose
        cols.add(col)
        diag.add(row - col)
        anti.add(row + col)
        found += place(row + 1, n, queens, cols, diag, anti)   # explore
        queens.pop()                   # un-choose
        cols.remove(col)
        diag.remove(row - col)
        anti.remove(row + col)
    return found


print(place(0, 4, [], set(), set(), set()), "solutions")
Output
.Q..
...Q
Q...
..Q.

..Q.
Q...
...Q
.Q..

2 solutions

How it works

  • When place is called for a row, every row above it already holds a queen that nothing attacks. So row == n means the board is full: one more solution.
  • Follow n = 4 from a queen in the top-left corner. The second row's queen can go in column 2 or 3. After column 2, every square of the third row is attacked. After column 3, the third row's queen fits in column 1, and then every square of the fourth row is attacked. Both ways are dead, so the search takes the first queen back and moves it to column 1, where it finds the first solution. Going back up like this, to change a choice made earlier, is what the word backtracking describes.
  • The choose changes four things, the list and the three sets, and the un-choose must undo every one of them. Forget one, say diag.remove, and that diagonal stays blocked for the rest of the search: the count comes out too small, with no error to warn you.

What it costs

Solutions multiply as the board grows: 352 for n = 9, 724 for n = 10, 2,680 for n = 11 and 14,200 for n = 12. The search is still exponential, but pruning cuts the tree down to a sliver: for n = 10 it makes 35,539 calls instead of checking 10¹⁰ boards. Each call does O(n) work, one O(1) check per column, and the recursion is only n + 1 calls deep, so the extra memory is O(n): the sets and the list of queens. Checking a new queen against every queen placed so far, with a loop, also works; it's O(n) per check instead of O(1), a few times slower. What makes the search fast is checking before placing, so that dead boards are never finished.

Your task

Challenge 1 of 3

A 4 × 4 board has queens at (0, 1) and (1, 3). What does this check of the third row print? Type each line it prints.

Challenge 2

Write count_queens(n), which returns the number of ways to put n queens on an n × n board so that no two attack each other. Don't print the boards. It must be fast for n = 11, so build the board row by row and skip attacked squares with the three sets.

For example, print(count_queens(4), count_queens(6)) prints:

2 4

Challenge 3

Often one solution is enough. Write first_solution(n): the first solution the search finds when it tries the columns from left to right, as a list of the queens' columns, row by row, or None when there's no solution. Stop as soon as you have one: make the search return True when the last row is placed, and return at once all the way up, without undoing anything. It must be fast for n = 21, which has hundreds of billions of solutions.

For example, first_solution(4) returns:

[1, 3, 0, 2]