[백준 9663] N-QUEEN(파이썬)

백준 9663 - NQUEEN

코드

import sys

N = int(sys.stdin.readline())

queen = [0] * N
flag_a = [False] * N #각 열 확인
flag_b = [False] * (2 * N - 1) #우상향 대각선 확인
flag_c = [False] * (2 * N - 1) #좌상향 대각선 확인
cnt = 0
answer = 0

def solution(N, i):
    global cnt
    for j in range(N):
        if (not flag_a[j] and not flag_b[i + j] and not flag_c[i - j + N - 1]):
            queen[i] = j
            if i == N - 1:
                cnt += 1
            else:
                flag_a[j] = flag_b[i + j] = flag_c[i - j + N - 1] = True
                solution(N, i + 1)
                flag_a[j] = flag_b[i + j] = flag_c[i - j + N - 1] = False
                
if N == 1:
    solution(N, 0)
    print(cnt)
elif N % 2 == 0:
    for k in range(N // 2):
        flag_a[k] = flag_b[k] = flag_c[- k + N - 1] = True
        solution(N, 1)
        flag_a[k] = flag_b[k] = flag_c[- k + N - 1] = False
        answer += cnt * 2
        cnt = 0
    print(answer)
else:
    for k in range(N // 2 + 1):
        flag_a[k] = flag_b[k] = flag_c[- k + N - 1] = True
        solution(N, 1)
        flag_a[k] = flag_b[k] = flag_c[- k + N - 1] = False
        if k != N // 2:
            answer += cnt * 2
            cnt = 0
        else:
            answer += cnt
    print(answer)

프로그래머스에서 풀어봤던 문제였지만 여기에서는 N의 최댓값이 15였기 때문에 pypy로는 통과할 수 있었지만 python으로 통과하기 위해서는 기존에 짰던 코드에서 시간을 굉장히 많이 줄여야 했다.

기존의 코드에서는

 def check(i, n):
     for j in range(i):
         if row[j] == n or abs(n - row[j]) == i - j:
             return False
     return True 

이 함수를 이용해 같은 열, 같은 대각선에 있는지를 확인하는 방법을 사용했는데 지금의 코드에서는 각 열, 대각선에 퀸이 놓여질 때 마다 flag_a, b, c에 체크해놓는 방법을 사용하여 연산 횟수를 많이 줄일 수 있었다.

마지막으로 모든 가능한 배치들이 좌우로 반전했을 때에도 가능한 배치가 된다는 점을 이용해 총 탐색 횟수를 반정도 줄일 수 있었다.

예를 들어 N = 5일 때 첫번째 열에 첫번째 칸에 퀸이 놓인 조건을 충족하는 배치들을 좌우로 반전했을 때 첫번째 열에 다섯번째 칸에 퀸이 놓인 조건을 충족하는 배치가 된다는 것이다.