shogi_cpsat.py

import itertools
from ortools.sat.python import cp_model

BOARD_SIZE = 9
WIDTH = HEIGHT = BOARD_SIZE
PIECETYPES = ['p', 'n', 's', 'g', 'k', 'l', 'r', 'b', 'q', '_']
# 0:up, 1:down, 2:left, 3:right, 4:upleft, 5:upright, 6:downleft, 7:downright,
# 8:knightleft, 9:knightright
DI = [-1, 1, 0, 0, -1, -1, 1, 1, -2, -2]
DJ = [0, 0, -1, 1, -1, 1, -1, 1, -1, 1]
PIECE_SHORTEFFECTS = {
    'p': [0],
    'n': [8, 9],
    's': [0, 4, 5, 6, 7],
    'g': [0, 1, 2, 3, 4, 5],
    'k': [0, 1, 2, 3, 4, 5, 6, 7],
    'l': [],
    'r': [],
    'b': [],
    'q': [],
    '_': [],
}
PIECE_LONGEFFECTS = {
    'p': [], 'n': [], 's': [], 'g': [], 'k': [],
    'l': [0], 'r': [0, 1, 2, 3], 'b': [4, 5, 6, 7], 'q': [0, 1, 2, 3, 4, 5, 6, 7],
    '_': [],
}
PIECE_NUM = {
    'p': 18, 'n': 4, 's': 4, 'g': 4, 'k': 2,
    'l': 4, 'r': 2, 'b': 2, 'q': 0, '_': 0
}

def getidx(i, j):
    return i * WIDTH + j
def is_inboard(i, j):
    return 0 <= i < HEIGHT and 0 <= j < WIDTH

class SolutionPrinter(cp_model.CpSolverSolutionCallback):
    """Print intermediate solutions."""

    def __init__(self, all_piece_board):
        cp_model.CpSolverSolutionCallback.__init__(self)
        self.__solution_count = 0
        self.__all_piece_board = all_piece_board

    def OnSolutionCallback(self):
        self.__solution_count += 1
        print(self.__solution_count)
        print('---')

        board = [[' ' for _j in range(WIDTH)] for _i in range(HEIGHT)]
        for pt, i, j in itertools.product(PIECETYPES, range(HEIGHT), range(WIDTH)):
            idx = getidx(i, j)
            v = self.__all_piece_board[pt][idx]
            if self.Value(v) == 1:
                board[i][j] = pt
        
        for i in range(HEIGHT):
            print(''.join(board[i]))
        print()


    def SolutionCount(self):
        return self.__solution_count

def main():
    model, all_piece_boards, exist_board = create_shogi_model()

    solver = cp_model.CpSolver()
    solution_printer = SolutionPrinter(all_piece_boards)
    status = solver.SearchForAllSolutions(model, solution_printer)
    print('t=', solver.WallTime())
    print('Solutions found :', solution_printer.SolutionCount())

def create_shogi_model():
    model = cp_model.CpModel()
    exist_board = [model.NewBoolVar('X_{}_{}'.format(i, j))
                   for i in range(HEIGHT) for j in range(WIDTH)]

    all_piece_boards = dict()
    for pt in PIECETYPES:
        pt_board = [model.NewBoolVar('{}_{}_{}'.format(pt, i, j))
                    for i in range(HEIGHT) for j in range(WIDTH)]
        model.Add(sum(pt_board) == PIECE_NUM[pt])
        all_piece_boards[pt] = pt_board

    for pt in PIECETYPES:
        pt_board = all_piece_boards[pt]
        for i, j in itertools.product(range(HEIGHT), range(WIDTH)):
            idx = getidx(i, j)
            model.AddImplication(pt_board[idx], exist_board[idx])

            for k in PIECE_SHORTEFFECTS[pt]:
                new_i = i + DI[k]
                new_j = j + DJ[k]
                if is_inboard(new_i, new_j):
                    model.AddImplication(pt_board[idx], exist_board[getidx(new_i, new_j)].Not())
                    # ないほうが高速に動作する
                    # for pt2 in PIECETYPES:
                    #     other_pt_board = all_piece_boards[pt2]
                    #     model.AddImplication(
                    #         pt_board[idx],
                    #         other_pt_board[getidx(new_i, new_j)].Not()
                    #     )

            for k in PIECE_LONGEFFECTS[pt]:
                new_i, new_j = i, j
                while True:
                    new_i += DI[k]
                    new_j += DJ[k]
                    if not is_inboard(new_i, new_j):
                        break
                    
                    model.AddImplication(pt_board[idx], exist_board[getidx(new_i, new_j)].Not())
                    for pt2 in PIECETYPES:
                        other_pt_board = all_piece_boards[pt2]
                        model.AddImplication(
                            pt_board[idx],
                            other_pt_board[getidx(new_i, new_j)].Not()
                        )


    for i, j in itertools.product(range(HEIGHT), range(WIDTH)):
        idx = getidx(i, j)

        sum_ij = sum([all_piece_boards[pt][idx] for pt in PIECETYPES])
        model.Add(sum_ij <= 1)

    return model, all_piece_boards, exist_board

if __name__ == "__main__":
    main()