leetcode---n-queens---dfs

The n-queens puzzle is the problem of placing n queens on an n×n chessboard such that no two queens attack each other.

Given an integer n, return all distinct solutions to the n-queens puzzle.

Each solution contains a distinct board configuration of the n-queens’ placement, where’Q’and’.’both indicate a queen and an empty space respectively.

For example,
There exist two distinct solutions to the 4-queens puzzle:

[
[“.Q..”, // Solution 1
“…Q”,
“Q…”,
“..Q.”],

[“..Q.”, // Solution 2
“Q…”,
“…Q”,
“.Q..”]
]

class Solution {
public:
    bool ok(int pos[], int row, int col)
    {
        for(int i=0; i<row; i++)
        {
            if(col == pos[i] || abs(i-row) == abs(pos[i]-col))
                return false;
        }
        return true;
    }

    void dfs(int row, int pos[], int n, vector<vector<string> > &ans)
    {
        if(row >= n)
        {
            vector<string> v(n, string(n, '.'));
            for(int i=0; i<n; i++)
            {
                v[i][pos[i]] = 'Q';
            }
            ans.push_back(v);
            return;
        }

        for(int i=0; i<n; i++)
        {
            pos[row] = i;
            if(ok(pos, row, i))
                dfs(row+1, pos, n, ans);
            pos[row] = -1;
        }
    }

    vector<vector<string> > solveNQueens(int n) 
    {
        vector<vector<string> > ans;
        int pos[n];
        dfs(0, pos, n, ans);
        return ans;
    }
};

Follow up for N-Queens problem.

Now, instead outputting board configurations, return the total number of distinct solutions.

class Solution {
public:

    bool ok(int pos[], int row, int col)
    {
        for(int i=0; i<row; i++)
            if(pos[i] == col || abs(i-row) == abs(pos[i]-col))
                return false;
        return true;
    }

    void dfs(int pos[], int row, int n, int &cnt)
    {
        if(row >= n)
            ++cnt;
        for(int i=0; i<n; i++)
        {
            pos[row] = i;
            if(ok(pos, row, i))
                dfs(pos, row+1, n, cnt);
            pos[row] = -1;
        }
    }

    int totalNQueens(int n) 
    {
        int pos[n];
        int cnt = 0;
        dfs(pos, 0, n, cnt);
        return cnt;
    }
};