1.7.4 N-Queens

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

/*

The N-queens problem asks for a placement of $n$ queens on an $n$-by-$n$ chessboard such that no two
share a row, column, or diagonal. Backtracking places one queen in each successive row and tracks
which columns and diagonals are occupied. A conflicting placement is rejected immediately; after a
failed recursive branch, its column and diagonals are released before trying the next column.

- `n_queens(n)` returns a placement as a vector `cols` of length `n`, where `cols[row]` is the
  queen's column in that row. It returns an empty vector if no placement exists.

Time Complexity:
- O(n!) per call.

Space Complexity:
- O(n) auxiliary and O(n) for the returned placement.

*/

#include <cassert>
#include <vector>

std::vector<int> n_queens(int n) {
  assert(n > 0);
  std::vector<int> cols(n);
  std::vector<char> used_col(n), used_diag1(2 * n - 1), used_diag2(2 * n - 1);
  auto rec = [&](auto &&rec, int row) -> bool {
    if (row == n) {
      return true;
    }
    for (int col = 0; col < n; col++) {
      int diag1 = row - col + n - 1, diag2 = row + col;
      if (used_col[col] || used_diag1[diag1] || used_diag2[diag2]) {
        continue;
      }
      cols[row] = col;
      used_col[col] = used_diag1[diag1] = used_diag2[diag2] = true;
      if (rec(rec, row + 1)) {
        return true;
      }
      used_col[col] = used_diag1[diag1] = used_diag2[diag2] = false;
      cols[row] = -1;
    }
    return false;
  };
  return rec(rec, 0) ? cols : std::vector<int>{};
}

/*** Example Usage ***/

#include <cassert>
#include <vector>
using namespace std;

int main() {
  assert(n_queens(4) == vector<int>({1, 3, 0, 2}));
  assert(n_queens(2).empty());
  assert(n_queens(1) == vector<int>({0}));
  return 0;
}