출처 | https://www.acmicpc.net/problem/9663

해당 블로그를 참고했다.

코드

n = int(input())

ans = 0
row = [0] * n

def promising(x):
    for i in range(x):
        if row[x] == row[i] or abs(row[x]-row[i]) == abs(x-i):
            return False
        
    return True


def n_queens(x):
    global ans
    if x == n:
        ans += 1
        return
    
    else:
        for i in range(n):
            row[x] = i
            if promising(x):
                n_queens(x+1)
                
n_queens(0)
print(ans)
        

코드 설명 및 풀이법

백트래킹의 가장 기본 문제 중 하나인 N-Queen 문제이다.

처음에는 퀸이 놓인 위치를 기록해주기 위해서 2차원 배열을 사용하려 했으나 단순히 1차원 배열의 인덱스와 값을 통해서 위치를 기록해줄 수 있었다.

row[i] = j는 다음과 같이 해석될 수 있다.

"퀸을 [i, j] 위치에 놓겠다."

퀸의 위치를 정하고 나서는, 해당 위치에 퀸을 놓을 수 있는지 없는지 is_promising 함수를 통해서 판단한다.

퀸을 놓지 못하는 경우는 2가지이다.

1) 같은 열에 다른 퀸이 있는 경우

이 경우는, row라는 배열 안에 같은 값이 있는지 없는지를 확인하면 된다.

즉, row[i] = j라고 정의한 것에서 동일한 j값이 있는지 보겠다는 것이다. 퀸의 위치가 기록되는 값은 [i, j]라고 보면 이해하기 쉽다.

2) 왼쪽 대각선, 오른쪽 대각선에 다른 퀸이 있는 경우.

우리는 퀸을 맨 윗 행부터 차례로 채우고 있다.

따라서 대각선을 확인할 때 현재 i보다 큰 값들은 확인할 필요 없이 i보다 작은 곳, 즉 체스판의 위쪽만 확인하면 된다.

대각선에 위치한 퀸을 확인하기 위해서는 그들 간 인덱스의 관계를 활용하면 O(1)번만에 확인할 수 있다.

첫 번째로, 왼쪽 위 대각선의 값들을 살펴보자.

만약에 현재 퀸을 놓은 위치가 (3, 3)라고 가정하면, 왼쪽 대각선의 좌표는 각각 (2, 2), (1, 1), (0, 0)이 된다.

여기서, (3, 3)을 i와 j라고 하고, (2, 2)를 x1, y1, (1, 1)을 x2, y2, (0, 0)을 x3, y3이라고 해보자.

i에서 x1을 뺀 값과 j에서 y1를 뺀 값은 모두 1로 같다.

또한, i에서 x2를 뺀 값과 j에서 y2를 뺀 값은 모두 2로 같다.

세 번째 경우도 3으로 마찬가지로 동일하다.

두 번째로, 오른쪽 위 대각선 값들을 살펴보자.

동일하게 현재 퀸의 위치가 (3, 3)이라고 가정하면 오른쪽 대각선의 좌표는 (2, 4), (1, 5), (0, 6)이 된다.

마찬가지로 (3, 3)을 i, j로, (2, 4)를 x1, y1, (1, 5)를 x2, y2, (0, 6)을 x3, y3로 두고 i, j값에서 x와 y를 뺀 값을 살펴보면

(1, -1), (2, -2), (3, -3)이 된다.

이 두 가지 케이스를 일반화하여 식으로 정리하면 is_promising 함수가 된다.

만약 현재 퀸을 놓으려는 위치가 promising하다면, 다음 퀸을 놓기 위한 n_queens(i + 1)을 호출해주도록 해주고, 그렇게 타고타고 들어간 x가 최종 깊이인 n이 되면 모든 퀸을 놓았다는 것이 되므로 ans에 1을 추가해준다.