Try to solve it yourself first, maybe using a simple brute force solution. Then check this post.

Problem:

In chess, queens can move any number of squares vertically, horizontally, or diagonally. The n-queens puzzle is the problem of placing n queens on an n × n chessboard so that no two queens can attack each other.

Given an integer n, print all possible distinct solutions to the n-queens puzzle. Each solution contains distinct board configurations of the placement of the n queens, where the solutions are arrays that contain permutations of [1, 2, 3, .. n]. The number in the $i^{th}$ position of the results array indicates that the $i^{th}$ column queen is placed in the row with that number. In your solution, the board configurations should be returned in lexicographical order.

Example:

For n = 1, the output should be nQueens(n) = [[1]].

For n = 4, the output should be

  nQueens(n) = [[2, 4, 1, 3],
                [3, 1, 4, 2]]

The N Queens problem is a very popular puzzle in computer science. Dijkstra provided a depth first search based backtracking solution to the problem in 1972, and sure enough it appears in the Backtracking category on CodeFights.

Backtracking is a problem solving technique that’s akin to brute force. What spares it from being brute force is that it prunes off subtrees in the search space based on rules defined by the specific problem being solved.

The nQueens problem is asking us to return an array of row numbers for each of the n queens placed in the n columns of the board. Note how we’re simplifying the problem from the get go by placing each queen in a separate column and returning just a $1d$ array of row numbers, rather than an $n \times n$ matrix. So in a way, we’ve manually pruned the search space ourselves before we even began writing code.

The remaining rules specific to this problem then are that no two row numbers should be the same, and no two queens should lie on the same diagonal.

N Queens

Checking the rows is easy. As for the diagonals, two $(row, col)$ pairs will lie on the same diagonal if the width of the triangle they form is equal to the height. That translates to,

Okay, now that we understand the conditions we need to test in order to prune off branches from the search tree, let’s dive straight into the code.

def nQueens(n):
  res = []
  for i in range(1, n+1):
    # TODO: add the recursive logic here
  return res

First off we’ll declare a list res to store the result, and at the end we’ll return it. The important stuff happens in the for loop. As I mentioned earlier, backtracking is almost brute force so we’re gonna consider each row from from $1$ to $n$ as our choice for placing the first queen (the queen in column 1).

Now comes the slightly tricky part. Having placed the first queen on some particular row, we’re going to use the nQueens function recursively to place the second queen. In general we’re going to need to know two things. First, we must know the placements of all the previous queens. Second, we’ll need to know the column we’re on, as a condition for stopping the recursion.

Let’s modify our code to include some of that.

def nQueens(n, state=[], col=1):
  if col > n: return [state]
  res = []
  for i in range(1, n+1):
    # TODO: add the recursive logic here
  return res

We could just compute col dynamically as col = len(state) + 1, but this is cleaner.

Each recursive call (which we’ll implement next) we make may return a list of solutions. We’ll have to make n such recursive calls for each of the n row positions for our first queen. We’ll simply collect these solutions in res and then return it. Notice how we’re updating the state in the call.

def nQueens(n, state=[], col=1):
  if col > n: return [state]
  res = []
  for i in range(1, n+1):
    for sol in nQueens(n, state + [i], col+1): res += [sol]
  return res

We’re only missing the critical piece now, which is the check to prevent an invalid configuration from ever being considered. This ensures that every solution returned is a valid one. Be reminded however that it may happen that some particular call to nQueens returns just an empty list. In this case the inner for loop never triggers and we just move on to the next row (i) on our outer for loop.

def nQueens(n, state=[], col=1):
  if col > n: return [state]
  res = []
  for i in range(1, n+1):
    for sol in nQueens(n, state + [i], col+1): res += [sol]
  return res

def invalid(s, r2):
  if not s: return False
  if r2 in s: return True
  c2 = len(s) + 1
  return any(abs(c1-c2) == abs(r1-r2) for c1, r1 in enumerate(s, 1))

Above is just the definition of the function invalid, but it’s not called yet. Let me first explain some of the details. The variables r1, c1, r2, c2 stand for row one, column one, row two, and column two. Row two, column two are the position of the potential square we are considering placing our queen on.

We first check if the state s is an empty list. This is a valid state since it just means we’re at the start and no queen has been placed on the board yet. Next we check if the row we’re planning on placing our queen already has another queen on it. If so, the configuration is invalid and we return True.

Finally, we check if $(r_2, c_2)$ shares a diagonal with any of the previous queens. We’re using the any() function with python’s list comprehension. The second argument to enumerate makes the count start at 1 instead of 0. With all that done, the only thing that remains is to use this function to trim down our search tree.

def nQueens(n, state=[], col=1):
  if col > n: return [state]
  res = []
  for i in range(1, n+1):
    if invalid(state, i): continue
    for sol in nQueens(n, state + [i], col+1): res += [sol]
  return res

def invalid(s, r2):
  if not s: return False
  if r2 in s: return True
  c2 = len(s) + 1
  return any(abs(c1-c2) == abs(r1-r2) for c1, r1 in enumerate(s, 1))

If our proposed configuration is invalid, we simply discard it and move on with the loop. That’s it, our final code.

Perhaps the hardest bit of the code to wrap your head around is the inner for loop. To recognize and master such recursive solutions two things are useful to have. First of course, is experience. The second, is to understand the art of wishful thinking. Wishful thinking is a bold thought process where you delay implementing all the intricacies, and minutia till the end and focus more on what you want from the recursive call, instead.

So the takeaway from this post is that for problems solved with backtracking, you can loop over all partial or complete solutions you get from the recursive call and then modify, and collect them as your complete solution to the problem.