52. N-Queens II

java

class Solution {
    int count = 0;
    public int totalNQueens(int n) {
        if (n <= 0) {
            return 0;
        }
        dfs(new boolean[n], new boolean[2 * n], new boolean[2 * n], 0, n);
        return count;
    }
    private void dfs(boolean[] shu,
                     boolean[] pie,
                     boolean[] na,
                     int row,
                     int n) {
        if (row == n) {
            count++;
            return;
        }
        for (int col = 0; col < n; col++) {
            int id1 = col + row;
            int id2 = n + col - row;
            if (shu[col] || pie[id1] || na[id2]) {
                continue;
            }
            shu[col] = true;
            pie[id1] = true;
            na[id2] = true;
            dfs(shu, pie, na, row + 1, n);
            na[id2] = false;
            pie[id1] = false;
            shu[col] = false;
        }
    }
}

python

class Solution:
    """
    @param: n: The number of queens.
    @return: The total number of distinct solutions.
    """
    count = 0
    def totalNQueens(self, n):
        # write your code here
        if n <= 0:
            return 0
        self.dfs([False] * n, [False] * 2 * n, [False] * 2 * n, 0, n)
        return self.count
        
    def dfs(self, shu, pie, na, row, n):
        if row == n:
            self.count += 1
            return
        for col in range(n):
            id1, id2 = col + row, n + col - row
            if shu[col] or pie[id1] or na[id2]:
                continue
            shu[col], pie[id1], na[id2] = True, True, True
            self.dfs(shu, pie, na, row + 1, n)
            shu[col], pie[id1], na[id2] = False, False, False