[백준] 9663 | N-Queen [파이썬/python]

문제

link: https://www.acmicpc.net/problem/9663

NxN 크기의 체스판 위에 N개의 queen이 서로를 공격하지 못하도록 위치시키는 경우의 수를 구하는 문제이다.

 

접근

알고리즘(백트래킹) 공부를 할 때 본 적 있는 문제라 반가웠다. 예전에 C++로 코드를 짜둔 게 있어서 python으로 바꿔 제출하려고 했는데 아래와 같이 시간 초과 오류가 났다. 코드 문제인 줄 알고 list도 set으로 바꾸고 이래저래 코드를 고쳐 제출했는데도 여전해서 난감했다...😅 결국 구글링했는데 PyPy3으로 제출해야 한다는 답을 얻었다! 더 이상의 코드 수정 없이 제출했다.

PyPy3으로 제출하니 해결된 모습

체스판의 queen은 행, 열, 대각선 방향으로 이동할 수 있다. set col에 queen N개가 위치하는 열의 숫자를 저장하고, 왼쪽 대각선에 위치하는지를 확인하기 위해 d_left에 행-열 값을, 오른쪽 대각선에 위치하는지를 확인하기 위해 d_right에 행+열 값을 저장하여 다음 queen의 위치가 유망한지를 확인하는 데 사용했다.

 

코드

import sys
input = sys.stdin.readline

N = int(input())

col = set()	#열
d_left = set()	#왼쪽 대각선
d_right = set()	#오른쪽 대각선
cnt = 0

def nqueens(i):
    global cnt
    
    #N개의 queen이 모두 놓였을 때
    if i == N:
        cnt += 1
        return
    
    for j in range(N):
    	#유망한지를 검사
        #이미 놓인 queen과 같은 열 or 왼쪽 대각선 or 오른쪽 대각선에 위치하는 경우
        if j in col or (i - j) in d_left or (i + j) in d_right:
            continue
        
        col.add(j)
        d_left.add(i - j)
        d_right.add(i + j)
        
        nqueens(i + 1)
        
        col.remove(j)
        d_left.remove(i - j)
        d_right.remove(i + j)
                
nqueens(0)
print(cnt)

 

+)

아래는 예전에 N-queen 문제를 해결하기 위해 작성한 C++ 코드이다. 경우의 수뿐만 아니라 queen의 위치도 출력한다.

#include <iostream>
#define MAX 50
using namespace std;

int col[MAX];
int n, total = 0;

bool promising(int i) {
	bool flag = true;
	for (int k = 1; k < i && flag; k++) {
		//Check if queens in row i and queens in row k are located in the same column or diagonal
		if (col[i] == col[k] || ((abs(col[i] - col[k]) == (i - k))))
			flag = false;
	}
	return flag;
}

void queens(int i) {
	if (promising(i)) {
		if (i == n) {
			cout << ++total << ": [";
			for(int k = 1; k <= n; k++)
				cout << "<" << k << ", " << col[k] << ">, ";
			cout << "\b\b]\n";
		}
		else {
			for (int j = 1; j <= n; j++) {
				col[i+1] = j;
				queens(i + 1);
			}
		}
	}
}

int main(void) {
  cout << "insert N: ";
	cin >> n;
	queens(0);
	cout << "Total number : " << total << "\n";
	system("pause");	//Preventing the console window from turning off immediately

	return 0;
}