PYTHON 2N皇后问题

问题描述
　　给定一个n*n的棋盘，棋盘中有一些位置不能放皇后。现在要向棋盘中放入n个黑皇后和n个白皇后，使任意的两个黑皇后都不在同一行、同一列或同一条对角线上，任意的两个白皇后都不在同一行、同一列或同一条对角线上。问总共有多少种放法？n小于等于8。
输入格式
　　输入的第一行为一个整数n，表示棋盘的大小。
　　接下来n行，每行n个0或1的整数，如果一个整数为1，表示对应的位置可以放皇后，如果一个整数为0，表示对应的位置不可以放皇后。
输出格式
　　输出一个整数，表示总共有多少种放法。
样例输入
4
1 1 1 1
1 1 1 1
1 1 1 1
1 1 1 1
样例输出
2
样例输入
4
1 0 1 1
1 1 1 1
1 1 1 1
1 1 1 1
样例输出
0


分析：和n皇后问题一样，只不过2n皇后要用两个深度优先搜索。
调用放置白皇后的递归dfs，先放置白皇后，当每一个白皇后放置成功之后，在递归的return语句之前，
新建一个棋盘，复制原来的棋盘后并把放置了白皇后的位置置为0，调用摆放黑皇后的深度优先搜索。

当黑皇后也找到相应的解法后，cnt++; 最后输出cnt的值。

import math

cnt = 0


def issafe(d, row):
    for i in range(row):
        if d[i] == d[row] or math.fabs(i - row) == math.fabs(d[i] - d[row]):
            return False
    return True


def blackDfs(c, d, n, row):
    global cnt
    if row == n:
        cnt += 1
        return
    d[row] = 0

    while d[row] < n:
        if c[row][d[row]] == 1 and issafe(c, d, row):
            blackDfs(c, d, n, row + 1)
        d[row] += 1


def dfs(a, b, n, row):
    if row == n:
        c = [[0 for i in range(n)]for i in range(n)]
        for i in range(n):
            for j in range(n):
                c[i][j] = a[i][j]  # save a to c

        for i in range(n):
            c[i][b[i]] = 0

        d = [0 for i in range(n)]
        blackDfs(c, d, n, 0)
        return

    b[row] = 0
    while b[row] < n:
        if a[row][b[row]] == 1 and issafe(b, row):
            dfs(a, b, n, row + 1)
        b[row] += 1


def main():
    n = int(input())

    a = [[0 for i in range(n)]for i in range(n)]
    b = [0 for i in range(n)]

    for i in range(n):
        for j in range(n):
            a[i][j] = int(input())
    dfs(a, b, n, 0)
    print(cnt)


if __name__ == '__main__':
    main()