[백준] 9663. N-Queen (Python)

📌 9663. N-Queen

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

풀이과정

2차원 배열을 사용하여 queen이 있는 경우 1로 표시해 두었다. 그리고 check 함수를 만들어서 상하좌우로 각각 가보며 1이 있는 경우를 찾고, 1이 있는 경우 False를 뱉게 만들었다.

import sys
input = sys.stdin.readline
n = int(input())

graph = [[0]*n for _ in range(n)]

def check(graph, a, b):
  dx = [1, -1, 0, 0, 1, -1, 1, -1]
  dy = [0, 0, 1, -1, -1, 1, 1, -1,]
  for i in range(8):
    nx, ny = a, b
    while nx >= 0 and nx<n and ny>=0 and ny<n:
      nx = nx + dx[i]
      ny = ny + dy[i]
      if nx >= 0 and nx<n and ny>=0 and ny<n and graph[nx][ny] == 1:
        return False
  return True

ans = 0
def n_queens(x):
  global ans
  if x == n:
    ans += 1
    return 
  else:
    for i in range(n):
      graph[x][i] = 1
      if check(graph, x, i):
        n_queens(x+1)
      graph[x][i] = 0

n_queens(0)
print(ans)

하지만 check에서 while이 많이 돌아서 시간초과 결과가 떴다. 그래서 while 대신 다른 방법으로 해볼 수 있나 찾아 보았는데, 대각선에 있는 경우에는 x값 차이의 절대값과 y값 차이의 절대값이 같다는 사실을 알게 되었다. 그래서 다음과 같이 check 함수를 바꾸었다.

def check(graph, a, b):
  for i in range(n):
    for j in range(n):
      if graph[i][j] == 1 and (i!=a or j!=b):
        if abs(a-i) == abs(b-j) or b == j:
          return False
  return True

하지만 이렇게 바꾸어도 시간초과가 여전히 나서, 이중 for문을 없애기 위해서 리스트에 queen이 있는 자리를 넣어 놓고, 그것에 대해서만 for문을 돌렸는데도 시간초과가 났다. 그 이유는 아직 잘 모르겠다. 만약, 이유를 알게 되면 나중에 수정을 하겠다.
다른 사람들의 풀이를 참고해보니, 1차원 배열로도 풀 수 있다는 것을 알게 되었다. 각 열에서 queen이 어느 인덱스에 있는지 인덱스값을 1차원 배열에 넣어주는 방식이다.

import sys
input = sys.stdin.readline

n = int(input())

ans = 0
row = [0] * n

def check(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 check(x):
       n_queens(x+1)

n_queens(0)
print(ans)

(그런데 백준 상의 문제인지, 이 코드로도 시간초과가 떴다. PyPy3로 제출하니 겨우 통과가 되었다.)

💯 정답

import sys
input = sys.stdin.readline

n = int(input())

ans = 0
row = [0] * n

def check(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 check(x):
        n_queens(x+1)

n_queens(0)
print(ans)