I've written a Python program to find the fastest way to a checkmate from a starting position. The stronger side has a King and two Bishops, and the weaker side has just a King. The weak King always tries to move towards the middle of the board, while the stronger side tries to checkmate.
I'm more looking for advice about the Python side of things than Chess-specific things. I'm particularly interested in knowing if the nested loops in the main function here could be refactored into something cleaner looking as the break
s and the state mutations make it hard for me to extract good functions out of them.
The code takes about 5 to 10 minutes to run, but it does reach and print an answer. The more specific the things that you say you would change, the better.
"""
What is the checkmate pattern for 1 king versus 1 king and 2 bishops?
This program finds and prints an example checkmate.
It works by doing a breadth first search from the start board, and breaking off the search
when a checkmate is found.
Throughout the program, the side with bishops is called the 'strong' side,
and the side without them the 'weak'.
"""
from collections import namedtuple
from typing import Set
Board = namedtuple('Board', 'weak_king, strong_king, bishop1, bishop2')
Square = namedtuple('Square', 'row col')
class BishopMovementWithStrongKing:
def can_get_to(self, start_square: Square, strong_king_square: Square) -> Set[Square]:
""" Where can a bishop get from a square, acknowledging that the strong king may block its movement?
We don't consider other possible blocking pieces than the strong king.
Calculate the result: slow"""
king_row, king_col = strong_king_square.row, strong_king_square.col
result = set()
for row_diff, col_diff in ((-1, -1), (-1, +1), (+1, -1), (+1, +1)):
temp_row, temp_col = start_square.row + row_diff, start_square.col + col_diff
while (0 <= temp_row <= 7) and (0 <= temp_col <= 7) and ((king_row, king_col) != (temp_row, temp_col)):
result.add(Square(temp_row, temp_col))
temp_row += row_diff
temp_col += col_diff
return result
class KingMovementIgnoreUnits:
""" Class to say what squares can a king can get to, ignoring restrictions caused by other units. """
def __init__(self):
self.destinations = [[self._calc_can_get_to(row, col) for col in range(8)] for row in range(8)]
def _calc_can_get_to(self, row, col) -> Set[Square]:
""" Where can a king get from the square at (row, col) ?
Calculate the result: slow"""
result = set()
for row_diff in (-1, 0, 1):
for col_diff in (-1, 0, 1):
if col_diff != 0 or row_diff != 0:
temp_row, temp_col = row + row_diff, col + col_diff # temp = destinations
if (0 <= temp_row <= 7) and (0 <= temp_col <= 7):
result.add(Square(temp_row, temp_col))
return result
def can_get_to(self, square: Square) -> Set[Square]:
""" Where can a king get from a square?
Lookup the result: fast"""
return self.destinations[square.row][square.col]
def are_adjacent(square1: Square, square2: Square) -> bool:
""" True if the two squares given by (col1, row1) and (col2, row2) are adjacent
ie they share an edge on a corner. """
return (abs(square1.col - square2.col) <= 1) and (abs(square1.row - square2.row) <= 1)
class KingMovementWithUnits:
"""
The squares a king can go, given that there are other units on the board,
and he can't enter check.
"""
def __init__(self):
self._kmiu = KingMovementIgnoreUnits()
self._bmwsk = BishopMovementWithStrongKing()
# How good (ie central) a square is for the losing king to move to.
# Lower is better.
self._square_rating = (
(61, 53, 45, 37, 38, 46, 54, 62),
(60, 33, 26, 17, 18, 27, 34, 55),
(52, 25, 13, 5, 6, 14, 28, 47),
(44, 23, 12, 1, 2, 7, 19, 39),
(43, 24, 11, 3, 4, 8, 20, 40),
(51, 32, 15, 9, 10, 16, 29, 48),
(59, 36, 30, 21, 22, 31, 35, 56),
(64, 58, 49, 41, 42, 50, 57, 63)
)
def weak_can_get_to(self, board: Board):
""" Return the set of squares the weak king can go to. """
return (self._kmiu.can_get_to(board.weak_king)
- self._kmiu.can_get_to(board.strong_king)
- self._bmwsk.can_get_to(board.bishop1, board.strong_king)
- self._bmwsk.can_get_to(board.bishop2, board.strong_king))
def strong_can_get_to(self, board: Board):
""" Return the set of squares the weak king can go to. """
return (self._kmiu.can_get_to(board.strong_king)
- self._kmiu.can_get_to(board.weak_king)
- {board.bishop1, board.bishop2})
def weak_goes(self, board: Board):
""" Return the best square available to the weak king, measured by its square rating,
or None if none are available. """
squares = self.weak_can_get_to(board)
if len(squares) == 0:
return None
# if weak king can take bishop 1, take it
elif are_adjacent(board.weak_king, board.bishop1) and not are_adjacent(board.bishop1, board.strong_king):
return board.bishop1
# if weak king can take bishop 2, take it
elif are_adjacent(board.weak_king, board.bishop2) and not are_adjacent(board.bishop2, board.strong_king):
return board.bishop2
else:
best_row, best_col = squares.pop()
best_score = self._square_rating[best_row][best_col]
for row, col in squares:
if self._square_rating[row][col] < best_score:
best_score = self._square_rating[row][col]
best_row, best_col = row, col
return Square(best_row, best_col)
def weak_to_move_boards(self, strong_to_move_board: Board):
""" Return the boards after a move from the strong side. """
result = set()
for new_king_sq in self.strong_can_get_to(strong_to_move_board):
result.add(strong_to_move_board._replace(strong_king=new_king_sq))
for new_bishop_sq in self._bmwsk.can_get_to(strong_to_move_board.bishop1, strong_to_move_board.strong_king):
result.add(strong_to_move_board._replace(bishop1=new_bishop_sq))
for new_bishop_sq in self._bmwsk.can_get_to(strong_to_move_board.bishop2, strong_to_move_board.strong_king):
result.add(strong_to_move_board._replace(bishop2=new_bishop_sq))
return result
class CheckCalculator():
""" For calculating if a king is in check. """
_bmwsk = BishopMovementWithStrongKing()
@classmethod
def weak_king_in_check(cls, board: Board) -> bool:
""" Return true if the weak king is in check and false otherwise. """
return board.weak_king in (cls._bmwsk.can_get_to(board.bishop1, board.strong_king) | cls._bmwsk.can_get_to(board.bishop2, board.strong_king))
def print_board(board: Board) -> None:
# letters meaning:
# w - weak king
# s - strong king
# b - bishop
letters = 'wsbb'
for row in range(8):
for col in range(8):
for piece_num in range(4):
if board[piece_num] == (row, col):
print(letters[piece_num], end=" ")
break
else:
if (row + col) % 2:
print("- ", end="")
else:
print("| ", end="")
print("")
print("")
def dict_chain_to_list(parent_by_child: dict, element):
""" Return a list from starting element to element with None as a parent. """
result = [element]
while element in parent_by_child.keys():
element = parent_by_child[element]
result.append(element)
return result
def main():
print("Looking for a checkmate")
kmwu = KingMovementWithUnits()
start_board = Board(weak_king=Square(0, 4), strong_king=Square(7, 4), bishop1=Square(7, 2), bishop2=Square(7, 5))
# In a round, these are the unexplored parent boards to create offspring boards from
unexplored_boards = {start_board, }
# These are the new boards found in a round
# to be explored in the next round
next_round_unexplored_boards = set()
# Parent boards of boards encountered previously
# Maps from child board to parent board, and
# ultimately will contain a tree from the start board
# outwards towards many future positions
# We map from one strong-side-to-move board to strong-side-to-move board here,
# skipping the weak to move boards.
seen_parents = {start_board: None}
# Maps from a board where the weak king is to move
# to the strong to move parent board
weak_to_move_parents = {}
checkmate_found = False
# Search for checkmate
while unexplored_boards and not checkmate_found: # loop once per round
# Search all known unexplored boards (except ones discovered in the process)
for parent_board in unexplored_boards: # explore the child positions of each new strong_to_move_board
# Find, record, and evaluate all children of a parent board
for weak_to_move_board in kmwu.weak_to_move_boards(parent_board):
# new_weak_king_pos will be a square if possible or None if there are no squares out of check
new_weak_king_pos = kmwu.weak_goes(weak_to_move_board)
# if weak king has somewhere to go ie not stalemate / checkmate
if new_weak_king_pos:
# if weak king has not taken a bishop
if new_weak_king_pos not in [weak_to_move_board.bishop1, weak_to_move_board.bishop2]:
child_board = weak_to_move_board._replace(weak_king=new_weak_king_pos)
if child_board not in seen_parents.keys():
seen_parents[child_board] = parent_board
weak_to_move_parents[child_board] = weak_to_move_board
next_round_unexplored_boards.add(child_board)
# Do nothing if weak king has taken bishop
else:
# Weak king can't move
checkmate_found = CheckCalculator.weak_king_in_check(weak_to_move_board)
if checkmate_found:
print("Found a checkmate")
checkmate_board = weak_to_move_board
checkmate_parent_board = parent_board
break
# Don't record stalemates
if checkmate_found:
break
if checkmate_found:
break
unexplored_boards = next_round_unexplored_boards
next_round_unexplored_boards = set()
if checkmate_found:
start_board_to_checkmate_parent = list(reversed(dict_chain_to_list(seen_parents, checkmate_parent_board)[:-1]))
weak_to_move_boards_first_to_checkmate = [weak_to_move_parents[strong_board]
for strong_board in start_board_to_checkmate_parent[1:]] + [checkmate_board]
# Print result
legend = '''Legend:
s strong king
w weak king
b bishop
| white square
- black square'''
print(legend)
print("")
for strong_to_move_board, weak_to_move_board in \
zip(start_board_to_checkmate_parent, weak_to_move_boards_first_to_checkmate):
print("Strong to move:")
print_board(strong_to_move_board)
print("Weak to move:")
print_board(weak_to_move_board)
else:
print("No checkmate found.")
main()