52. N-Queens II

Description

The n-queens puzzle is the problem of placing n queens on an n x n chessboard such that no two queens attack each other.

Given an integer n, return the number of distinct solutions to the n-queens puzzle.

 

Example 1:

Input: n = 4
Output: 2
Explanation: There are two distinct solutions to the 4-queens puzzle as shown.

Example 2:

Input: n = 1
Output: 1

 

Constraints:

• 1 <= n <= 9

Solution

n-queens-ii.py

class Solution:
    def totalNQueens(self, n: int) -> int:
        cols, diag1, diag2 = set(), set(), set()

        def go(row):
            if row == n: return 1

            res = 0

            for col in range(n):
                if col not in cols and row - col not in diag1 and row + col not in diag2:
                    cols.add(col)
                    diag1.add(row - col)
                    diag2.add(row + col)
                    res += go(row + 1)
                    cols.remove(col)
                    diag1.remove(row - col)
                    diag2.remove(row + col)

            return res

        return go(0)