2n皇后问题（python）

采用深度优先+剪枝的算法。
基本思路是先排序黑皇后，再排序白皇后。

n = int(input())
ls = [[int(x) for x in input().split()] for i in range(n)]
# 结果res
res = 0
# 代表hasblack[j]第j列是否有black,1代表有
hasblack = [0 for i in range(n)]
haswhite = [0 for i in range(n)]
# 递归的深度优先加剪枝

# 剪枝函数
def isOk(i, j, bw):
    # 考虑同列是否有
    if bw == 2 and hasblack[j] == 1:
        return False
    if bw == 3 and haswhite[j] == 1:
        return False
    # 考虑左上面的对角线
    t = min(i, j)
    for k in range(t):
        if ls[i-k-1][j-k-1] == bw:
            return False
    # 考虑右上方对角线
    t = min(i, n-j-1)
    for k in range(t):
        if ls[i-k-1][j+k+1] == bw:
            return False
    return True

# 在第i行放black or white皇后
def dfs(ls, i, bw):
    global res, hasblack, haswhite
    # 递归出口
    if i == n:
        # 白皇后放完了，res ++
        if bw == 3:
            res += 1
        # 黑皇后放完了，放白
        else:
            dfs(ls, 0, 3)
            
    else:
        # 放黑皇后
        if bw == 2:
            for j in range(n):
                # 能放且不与之前的冲突
                if ls[i][j] == 1 and isOk(i, j, bw):
                    ls[i][j] = bw
                    hasblack[j] = 1
                    dfs(ls, i+1, bw)
                    hasblack[j] = 0
                    ls[i][j] = 1
        #放白皇后
        if bw == 3:
            for j in range(n):
                # 能放且不与之前的冲突
                if ls[i][j] == 1 and isOk(i, j, bw):
                    ls[i][j] = bw
                    haswhite[j] = 1
                    dfs(ls, i+1, bw)
                    haswhite[j] = 0
                    ls[i][j] = 1


            
dfs(ls, 0, 2)
print(res)