Grokking 194

class Solution {

  List<List<String>> solutions;

  Set<Integer> cols;
  Set<Integer> diagonals; // represent diagonal indices
  Set<Integer> antiDiagonals; //represent anti diagonal indices
  int size = 0;

  public List<List<String>> solveNQueens(int n) {
    size = n;
    solutions = new ArrayList<List<String>>();
    cols = new HashSet<Integer>();
    diagonals = new HashSet<Integer>();
    antiDiagonals = new HashSet<Integer>();

    // Create empty board
    char board[][] = new char[size][size];
    for (int row = 0; row < n; row++) {
      for (int col = 0; col < n; col++) {
        board[row][col] = '.';
      }
    }

    backtrack(board, 0);

    return solutions;
  }

  private void backtrack(char[][] board, int row) {
    if (row == size) {
      solutions.add(createBoard(board));
      return;
    }

    for (int col = 0; col < size; col++) {
      int currDiagonal = row - col;
      int currAntiDiagonal = row + col;

      // If queen can not be placed, try next column
      if (cols.contains(col) || diagonals.contains(currDiagonal) || antiDiagonals.contains(currAntiDiagonal)) {
        continue;
      }

      // Place a queen
      board[row][col] = 'Q';
      cols.add(col);
      diagonals.add(currDiagonal);
      antiDiagonals.add(currAntiDiagonal);

      // Try to place a queen at new row
      backtrack(board, row + 1);

      // Remove placed queen
      board[row][col] = '.';
      cols.remove(col);
      diagonals.remove(currDiagonal);
      antiDiagonals.remove(currAntiDiagonal);
    }
  }

  private List<String> createBoard(char[][] board) {
    List<String> res = new ArrayList<String>();

    for (int row = 0; row < size; row++) {
      String current_row = new String(board[row]);
      res.add(current_row);
    }

    return res;
  }