Skip to content

Instantly share code, notes, and snippets.

@komori-n
Last active Nov 24, 2019
Embed
What would you like to do?
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment