N-Queens II LeetCode solution in Python

1
class Solution:
2
    def totalNQueens(self, n: int) -> int:
3
        res = 0
4
        # 用于存放结果
5
        pdia = set()
6
        ndia = set()
7
        col = set()
8

9
        def backtrack(r):
10
            # 利用r作为一种计数，表示目前所在的行数
11
            if r == n:
12
                # 判断已经完成棋盘，返回结果
13
                nonlocal res
14
                res += 1
15
                return
16
            for c in range(n):
17
                # 对于n行n列的棋盘，每次在每一行我们尝试n种选择，
18
                # 即每个岔路口有n条路线
19
                if (r + c) in pdia or (r - c) in ndia or c in col:
20
                    # 如果在同一对角线，或同一列，则不符合要求
21
                    continue
22
                col.add(c)
23
                pdia.add(r + c)
24
                ndia.add(r - c)
25

26
                backtrack(r + 1)
27

28
                col.remove(c)
29
                pdia.remove(r + c)
30
                ndia.remove(r - c)
31

32
        backtrack(0)
33
        return res