【DFS】蓝桥杯试题 基础练习 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

数据规模和约定
　　我们约定，输入的串的长度不超过15
　　资源约定：峰值内存消耗（含虚拟机） < 256M， CPU消耗 < 1000ms

思路

先放黑子、再放子
1、在原始输入棋盘的基础上，先用一个vector保存所有放了n个黑子的合法键盘
2、再用另一个vector保存在每一个放了n个黑子的合法键盘上放n个白子的合法键盘，该vector的大小就是答案

代码

#include<bits/stdc++.h>
#define N 10
using namespace std;
vector<vector<int> > a(N, vector<int>(N));  //放原始棋盘
vector<vector<int> > b(N, vector<int>(N, 1));   //中间过渡棋盘
vector<vector<vector<int> > > all;  //存储所有放n个黑子的合法棋盘
vector<vector<vector<int> > > ans;  //储存在all的棋盘中放n个白子的合法棋盘
int n;

//放子后更新棋盘
void update(vector<vector<int> > &arr, int i, int j, int val){
    //同行
    for(int k = i+1; k<n; k++)
        arr[k][j] = val;

    //同列
    for(int k=j+1; k<n; k++)
        arr[i][k] = val;

    //正对角线
    for(int k=1; i+k<n&&j+k<n; k++)
        arr[i+k][j+k] = val;

    //负对角线
    for(int k=1; i+k<n&&j-k>=0; k++)
        arr[i+k][j-k] = val;
}

//恢复棋盘
void back(vector<vector<int> > &arr, vector<vector<int> > &c, int i, int j){

    for(int k = i+1; k<n; k++)
        arr[k][j] = c[k][j];

    for(int k=j+1; k<n; k++)
        arr[i][k] = c[i][k];

    for(int k=1; i+k<n&&j+k<n; k++)
        arr[i+k][j+k] = c[i+k][j+k];

    for(int k=1; i+k<n&&j-k>=0; k++)
        arr[i+k][j-k] = c[i+k][j-k];
}

//放黑子
void putBlack(int k){
    //放完n个黑子，保存棋盘
    if(k==n){
        all.push_back(b);
        return;
    }

    for(int i=0; i<n; i++){
        if(a[k][i]==0 || b[k][i]==0)
            continue;
        if(k==0)
            b.assign(a.begin(), a.end());
        //用多一个vector保存棋盘，以便后面恢复棋盘
        vector<vector<int> > c(b.begin(), b.end());

        //放黑子
        b[k][i] = 2;

        //更新棋盘
        update(b, k, i, 0);

        //继续在第k+1行放黑子
        putBlack(k+1);

        //恢复棋盘
        back(b, c, k, i);
        b[k][i] = 1;
    }
}

//放白子
void putWhite(vector<vector<int> > &arr, int k){
    //放完n个白子 保存棋盘
    if(k==n){
        ans.push_back(b);
        return;
    }

    for(int i=0; i<n; i++){
        if(a[k][i]==0 || arr[k][i]==2 || b[k][i]==5)
            continue;
        if(k==0)
            b.assign(arr.begin(), arr.end());

        //用多一个vector保存棋盘，以便后面恢复棋盘
        vector<vector<int> > c(b.begin(), b.end());

        //放白子
        b[k][i] = 3;

        //更新棋盘
        update(b, k, i, 5);

        //继续在第k+1行放白子
        putWhite(arr, k+1);

        //恢复棋盘
        back(b, c, k, i);
        b[k][i] = 1;
    }
}

int main(){
    cin>>n;
    for(int i=0; i<n; i++)
        for(int j=0; j<n; j++)
            cin>>a[i][j];

    //从第0行放黑子
    putBlack(0);

    //初始化b
    for(int i=0; i<n; i++)
        for(int j=0; j<n; j++)
            b[i][j] = 1;

	//放白子
    for(int i=0; i<all.size(); i++){
        //初始化b
        for(int j=0; j<n; j++){
            for(int k=0; k<n; k++){
                b[j][k] = 1;
            }
        }
        //从all中每个棋盘的第0行放白子
        putWhite(all[i], 0);
    }

    //打印所有棋盘个数
    cout<<ans.size()<<endl;

    return 0;
}