Algorithm/LeetCode

[LeetCode # 51] N-Queens

꼽냥이 2021. 6. 18. 20:24
반응형

체스판 위에 N 개의 Queen 을 서로 공격할 수 없도록 놓는 문제

 

유명한 BackTracking 문제라 DFS 를 활용해 문제를 풀이했다. 문제의 답으로 체스판을 문자열로 표현해야 해서 빈 체스판을 먼저 저장해 놓고, Queen 을 한 개씩 배치하며 해당 위치에 배치 가능한 지 확인했다. 총 N 개 Queen 이 배치된 후 체스판을 결과물 저장하는 Vector 에 추가했다.

 

이러한 풀이로 작성한 코드는 다음과 같다.

 

class Solution {
public:
    vector<vector<string>> ans;
    int dr[4] = {-1, -1, 1, 1}, dc[4] = {-1, 1, -1, 1};
    
    bool check(vector<string>& curVec, int i, int j, int n) {
        // 가로 세로
        for (int k = 0; k<n; k++) {
            if (curVec[i][k] == 'Q')
                return false;
            if (curVec[k][j] == 'Q')
                return false;
        }
        // 대각선
        for (int k = 1; ; k++) {
            bool flag = false;
            for (int l = 0; l < 4; l++) {
                int ni = i + dr[l] * k, nj = j + dc[l] * k;
                if (ni >= 0 && ni < n && nj >= 0 && nj < n) {
                    flag = true;
                    if (curVec[ni][nj] == 'Q')
                        return false;
                }
            }
            if (!flag)
                break;
        }
        return true;
    }
    void solveRecursive(vector<string>& curVec, int r, int cnt, int n) {
        if (cnt == n) {
            ans.push_back(curVec);
            return;
        }

        for (int i = r; i<n; i++) {
            for (int j = 0; j<n; j++) {
                if (curVec[i][j] == '.' && check(curVec, i, j, n)) {
                    curVec[i][j] = 'Q';
                    solveRecursive(curVec, i + 1, cnt + 1, n);
                    curVec[i][j] = '.';
                }
            }
        }
    }
    vector<vector<string>> solveNQueens(int n) {
        string s;
        for (int i = 0; i<n; i++) {
            s += ".";
        }
        vector<string> vec;
        for (int i = 0; i<n; i++) {
            vec.push_back(s);
        }
        solveRecursive(vec, 0, 0, n);
        return ans;
    }
};

 

위와 같이 코드를 작성했더니 시간이 너무 오래 걸렸다. 불필요한 조건 검사가 많고 생각을 너무 단순하게 해서 그런 것으로 생각이 되었다. 우선, 위에서 가로 검사의 경우는 한 줄에 하나의 Queen 만 배치가 가능하므로 굳이 검사를 할 필요가 없고, 세로 검사의 경우도 현재 위치보다 아래 줄은 검사를 할 필요가 없다. 또한 대각선 검사의 경우도 마찬가지로 현재 위치보다 윗줄만 검사를 하면 되어 이렇게 수정해 코드를 다시 제출했다.

 

class Solution {
public:
    vector<vector<string>> solveNQueens(int n) {
        vector<vector<string>> ans;
        vector<string> nQueen(n, string(n, '.'));
        solveNQueens(ans, nQueen, 0, n);
        return ans;
    }
    void solveNQueens(vector<vector<string>>& ans, vector<string>& nQueen, int r, int n) {
        if (r == n) {
            ans.push_back(nQueen);
            return;
        }

        for (int j = 0; j<n; j++) {
            if (check(nQueen, r, j)) {
                nQueen[r][j] = 'Q';
                solveNQueens(ans, nQueen, r + 1, n);
                nQueen[r][j] = '.';
            }
        }
    }
    bool check(vector<string>& nQueen, int i, int j) {
        for (int k = 1; i - k >= 0; k++) {
            if (nQueen[i - k][j] == 'Q')
                return false;
            if (j - k >= 0 && nQueen[i - k][j - k] == 'Q')
                return false;
            if (j + k < nQueen.size() && nQueen[i - k][j + k] == 'Q')
                return false;
        }
        return true;
    }
};

 

이렇게 풀이를 변경한 후 전반적인 소스 코드도 더 간결해졌고, 시간도 훨씬 적게 걸렸다.