#2 OrTools로 n-Queens 솔루션 결과 출력

이전 글에서 orTools로 n-Queens 솔루션을 찾는 방법을 설명했습니다. 이번 글에서는 지난 번 글에 이어지는 글로 솔루션의 결과를 출력하는 방법을 소개합니다. 

  ### 모델 해결
  solver = cp_model.CpSolver()
  ### 모델 다이어그램 출력
  printer = DiagramPrinter(queens)
  ### 모델 다이어그램 출력
  status = solver.SearchForAllSolutions(model, printer)
  print()
  print('찾은 솔루션 개수 : %i' % printer.SolutionCount())
  print("걸린 시간: ", solver.WallTime(), "ms")

지난 번 소스 맨 마지막 부분을 보면 solver를 통해 솔루션을 찾기에서  DiagramPrinter라는 클래스 인스턴스를 생성하고 SearchForAllSolutions에 인자로 전달하는 부분이 있습니다. 해당 클래스는 솔루션 결과에 대한 후속 작업을 정의할 수 있는 콜백 클래스라고 보시면 됩니다. 아래 DiagramPrinter 클래스 정의 소스를 보시면 __init__, OnSolutionCallback, SolutionCount 라는 메소드를 정의하고 있습니다.

class DiagramPrinter(cp_model.CpSolverSolutionCallback):
  def __init__(self, variables):
    cp_model.CpSolverSolutionCallback.__init__(self)
    self.__variables = variables
    self.__solution_count = 0

  def OnSolutionCallback(self):
    self.__solution_count += 1

    for v in self.__variables:
      queen_col = int(self.Value(v))
      board_size = len(self.__variables)
      # 퀸이 있는 행 인쇄
      for j in range(board_size):
        if j == queen_col:
          # j열, i행에 여왕이 있으면 Q 인쇄
          print("Q", end=" ")
        else:
          print("_", end=" ")
      print()
    print()

  def SolutionCount(self):
    return self.__solution_count

__init__메소드를 통해 초기화하고, OnSolutionCallback이란 메소드를 통해 솔루션이 있을 때마다 솔루션 count를 1증가하고 화면에 queens 배열에 저장된 솔루션 값을 인쇄합니다.

전체 소스는 아래와 같습니다.

 

import sys
# CP-SAT 솔버
from ortools.sat.python import cp_model

def main(board_size):
  model = cp_model.CpModel()
  # 변수 생성
  # 배열의 인덱스는 열(column), 값은 행(Row)
  queens = [model.NewIntVar(0, board_size - 1, 'x%i' % i) for i in range(board_size)]
  
  # 제약조건 생성
  # 모든 여왕은 다른 행에 있어야 한다는 조건
  model.AddAllDifferent(queens)
  # 참고: 배열은 인덱스가 모두 다르므로 기본적으로 모든 여왕은 다른 열에 있고,
  #       배열의 값을 모두 다르게 지정하면 여왕은 모두 다른 행에 위치 하게 됨

  # 두 개의 여왕이 동일한 대각선에 있을 수 없다는 제약 조건 설정
  diag1 = []
  for i in range(board_size):
    q1 = model.NewIntVar(0, 2 * board_size, 'diag1_%i' % i)
    diag1.append(q1)
    model.Add(q1 == queens[i] + i)
      
  diag2 = []
  for i in range(board_size):
    q2 = model.NewIntVar(-board_size, board_size, 'diag2_%i' % i)
    diag2.append(q2)
    model.Add(q2 == queens[i] - i)

  model.AddAllDifferent(diag1)
  model.AddAllDifferent(diag2)
  
  ### 모델 해결
  solver = cp_model.CpSolver()
  ### 모델 다이어그램 출력
  printer = DiagramPrinter(queens)
  ### 모델 다이어그램 출력
  status = solver.SearchForAllSolutions(model, printer)
  print()
  print('찾은 솔루션 개수 : %i' % printer.SolutionCount())
  print("걸린 시간: ", solver.WallTime(), "ms")


class DiagramPrinter(cp_model.CpSolverSolutionCallback):
  def __init__(self, variables):
    cp_model.CpSolverSolutionCallback.__init__(self)
    self.__variables = variables
    self.__solution_count = 0

  def OnSolutionCallback(self):
    self.__solution_count += 1

    for v in self.__variables:
      queen_col = int(self.Value(v))
      board_size = len(self.__variables)
      # 퀸이 있는 행 인쇄
      for j in range(board_size):
        if j == queen_col:
          # j열, i행에 여왕이 있으면 Q 인쇄
          print("Q", end=" ")
        else:
          print("_", end=" ")
      print()
    print()

  def SolutionCount(self):
    return self.__solution_count

board_size = 4
main(board_size)

참고로 python 문법으로 소스를 좀더 축약 시킬 수 있는 부분이 있습니다. 소스 18-28행은 다음과 같이 단 2줄로 축약이 가능합니다.

diag1 = [queens[i] + i for i in range(board_size)]
diag2 = [queens[i] - i for i in range(board_size)]