蓝桥杯 基础练习 2n皇后（Python）

资源限制

内存限制：512.0MB   C/C++时间限制：1.0s   Java时间限制：3.0s   Python时间限制：5.0s

问题描述

　　给定一个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

from itertools import *

ss = ['1', '2', '3', '4', '5', '6', '7', '8']
n = int(input())
s = ss[:n]
nums = []  # 棋盘
for i in range(n):
    nums.append(list(map(int, input().split())))
queens = []  # 皇后放置位置
for i in permutations(s):  # 列出所有皇后的存放方式（此时按一行放一个）（包括黑白）
    a = ''.join(i)
    queens.append(a)
for i in range(len(queens) - 1, -1, -1):  # 剔除不符合题意的放置方案
    s1 = queens[i]
    tmp = 0
    for j in range(n - 1):
        for k in range(1, n - j):
            if abs(int(s1[j]) - int(s1[j + k])) == k:
                queens.pop(i)
                tmp = 1
            if tmp == 1:
                break
        if tmp == 1:
            break
double_queens = []
for i in permutations(queens, 2):  # 列出所有黑白皇后的放置方案
    a = ''.join(i)
    double_queens.append(a)
for i in range(len(double_queens) - 1, -1, -1):  # 删除有重复位置的方案
    tmp = 0
    s1 = double_queens[i]
    for j in range(n):
        if s1[j] == s1[j + n]:
            double_queens.pop(i)
            tmp = 1
        if tmp == 1:
            break
m = 0
for i in double_queens:  # 跳过占用棋盘中0位置的方案，并且记录符合条件的方案数
    tmp = 0
    for j in range(n):
        if nums[j][int(i[j]) - 1] == 0 or nums[j][int(i[j + n]) - 1] == 0:
            tmp = 1
        if tmp == 1:
            break
    if tmp == 1:
        continue
    elif tmp == 0:
        m += 1
print(m)