Finding a checkmate with two bishops using Python

Question

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 breaks 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()

Answer 1

Style

Other than some overly long lines, there is not much that stand out.

Some names feels a bit confusing at first like weak_to_move which is after a move from the strong side; and, generaly, the whole namming scheme that have too much names using the strong/weak wording makes it hard to follow.

The unexplored_boards = {start_board, } can be simplified to unexplored_boards = {start_board} and I'd use a constant like BOARD_LIMITS = range(8) to simplify understanding some bits of code:

• printing a board can use for row in BOARD_LIMITS: for col in BOARD_LIMITS: ...;
• checking that a position is inside the board can be turned into if row in BOARD_LIMITS and col in BOARD_LIMITS

Lastly, you should use an if __name__ == '__main__': clause to avoid running the whole search each time you import your file for testing purposes. Speaking of which, the main function may be parametized by the starting board so you can easily test from several starting positions.

Checkmates and replaying the game

Instead of having your main() function do two things: search for a checkmate and print the boards that lead to that position, I would split the behaviour into more functions. Use a function that perform the search (basically only the while unexplored_boards and a bit of initialization) and call it from main. This way, you would be able to return from this function without having to use these awkward multiple if checkmate_found.

I would also change dict_chain_to_list to simplify the amount of work needed by main: make this function take both the strong parents and the weak parents to be able to generate each position on its own. You could also turn it into a generator to avoid creating yet another list; but most importantly, you should make it so that the position are generated in order from start to finish.

Lastly, I would make the (newly created) search_for_checkmate function return only the generator created by such function (plus the last board where the weak side is checkmate which is in neither dictionary); so that, again, the amount of work needed by main is reduced.

Printing boards

You can simplify the print_board function if you remember that strings and namedtuples are both iterable. So you can zip(board, letters). And if you feed that to a dict constructor, you have a mapping from Squares to the associated letter.

Except the blank squares are missing which can easily be fixed by using the dict.get method with a default value of '|–'[(row + col) % 2].

Using such approach, you can easily come up with:

def print_board(board):
    hot_positions = dict(zip(board, 'wsbb'))
    board_string = '\n'.join(
            ' '.join(
                hot_positions.get((row, col), '|—'[(row + col) % 2])
                for col in BOARD_LIMITS)
            for row in BOARD_LIMITS)
    print(board_string)

Which is nice but could also be a __str__ method on the Board class so you could just print(board) to start with:

class Board(namedtuple('Board', 'weak_king strong_king bishop1 bishop2')):
    def __str__(self):
        hot_positions = dict(zip(self, 'wsbb'))
        return '\n'.join(
                ' '.join(
                    hot_positions.get((row, col), '|—'[(row + col) % 2])
                    for col in BOARD_LIMITS)
                for row in BOARD_LIMITS)

OOP

Now we're getting into the core of what I feel is bad design. You have too much class used for a single method. Possibly the constructor performs some kind of caching beforehand. This make the code unnecessarily complex: use simple functions instead.

Or better, since most of these functions take a board as parameter, make them method of the Board class, it's more logical. In fact, appart from the special main, search_for_checkmate and dict_chain_to_list functions, pretty much everything else can be seen as a method of either Square or Board:

• generating possible positions for a piece from a given Square;
• checking that a given Square is valid within the chess board;
• moving a Square with a given offset (which is just an helper for the positions generators);
• generating the possible Boards after a move from the strong side;
• generating the possible Boards after a move from the weak side;
• selecting the best Board for a move of the weak side;
• checking that the Board is missing a bishop;
• checking that the weak king on the Board is in check.

Make Square as generic as possible and generate positions regardless of where other pieces are. And filter those position in Board as necessary.

Proposed improvements

"""Two bishops to checkmate

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'.
"""

import itertools
from collections import namedtuple
from typing import Set


BOARD_LIMITS = range(8)
# How good (ie central) a square is for the weak king to move to.
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)
)


class Square(namedtuple('Square', 'row col')):
    def move(self, row, col):
        return Square(self.row + row, self.col + col)

    @property
    def is_valid(self):
        return all(x in BOARD_LIMITS for x in self)

    def king_movements(self):
        for row, col in itertools.product(range(-1, 2), repeat=2):
            new_position = self.move(row, col)
            if new_position != self and new_position.is_valid:
                yield new_position

    def bishop_movements(self):
        # First diagonal
        offset = min(self)
        yield from self._bishop_helper(-offset, 1)

        # Second diagonal
        offset = min(self.row, 7 - self.col)
        yield from self._bishop_helper(-offset, -1)

    def _bishop_helper(self, offset, column_adjustment):
        position = self.move(offset, column_adjustment * offset)
        while position.is_valid:
            if position != self:
                yield position
            position = position.move(1, column_adjustment)


class Board(namedtuple('Board', 'weak_king strong_king bishop1 bishop2')):
    def strong_moves(self):
        king_restricted_moves = set(self.weak_king.king_movements())
        king_restricted_moves.add(self.bishop1)
        king_restricted_moves.add(self.bishop2)
        for king_move in self.strong_king.king_movements():
            if king_move not in king_restricted_moves:
                yield Board(self.weak_king, king_move, self.bishop1, self.bishop2)
        for bishop_move in self.bishop1.bishop_movements():
            if self._bishop_can_move(self.bishop1, bishop_move):
                yield Board(self.weak_king, self.strong_king, bishop_move, self.bishop2)
        for bishop_move in self.bishop2.bishop_movements():
            if self._bishop_can_move(self.bishop2, bishop_move):
                yield Board(self.weak_king, self.strong_king, self.bishop1, bishop_move)

    def weak_moves(self):
        king_restricted_moves = set(self.strong_king.king_movements())
        king_restricted_moves.update(self.bishop1.bishop_movements())
        king_restricted_moves.update(self.bishop2.bishop_movements())
        for king_move in self.weak_king.king_movements():
            if king_move not in king_restricted_moves:
                yield Board(king_move, self.strong_king, self.bishop1, self.bishop2)

    def _bishop_can_move(self, bishop, destination):
        x1, y1 = destination
        x2, y2 = bishop
        intermediate_coordinates = zip(range(x1, x2, sign(x2 - x1)), range(y1, y2, sign(y2 - y1)))
        possibilities = {Square(x, y) for x, y in intermediate_coordinates}
        return not (set(self) & possibilities)

    @property
    def is_bishop_taken(self):
        return self.weak_king == self.bishop1 or self.weak_king == self.bishop2

    @property
    def best_weak_move(self):
        return min(self.weak_moves(), key=rate_board)

    @property
    def weak_king_in_check(self):
        king = self.weak_king
        return (
            any(bishop == king for bishop in self.bishop1.bishop_movements()) or
            any(bishop == king for bishop in self.bishop2.bishop_movements())
        )

    def __str__(self):
        hot_positions = dict(zip(self, 'wsbb'))
        return '\n'.join(
                ' '.join(
                    hot_positions.get((row, col), '|—'[(row + col) % 2])
                    for col in BOARD_LIMITS)
                for row in BOARD_LIMITS)


def rate_board(board):
    """Rate a board for the weak side. Lower is better."""
    if board.is_bishop_taken:
        return 0
    row, col = board.weak_king
    return SQUARE_RATING[row][col]


def sign(x):
    return x // abs(x)


def rewind_game(strong_to_move_parents: dict, weak_to_move_parents: dict, element: Board):
    """Generate boards from starting element to the given one.

    The given element is assumed to be a strong-to-move board.
    """

    previous = strong_to_move_parents[element]
    if previous is not None:
        # Start board not reached yet
        yield from rewind_game(strong_to_move_parents, weak_to_move_parents, previous)
        yield weak_to_move_parents[element]
    yield element


def search_for_checkmate(start_board: Board):
    """Explore board positions using a breadth first search approach
    to find a checkmate starting at the given board.

    Return an iterable of boards leading to the found checkmate and
    the board of said checkmate.
    """

    unexplored = {start_board}
    strong_parents = {start_board: None}
    weak_parents = {}

    while unexplored:
        next_round_unexplored = set()

        for strong_to_move in unexplored:
            # Find, record, and evaluate all children of a parent board
            for weak_to_move in strong_to_move.strong_moves():
                try:
                    new_weak_king_position = weak_to_move.best_weak_move
                except ValueError:
                    # Weak king can't move
                    if weak_to_move.weak_king_in_check:
                        game = rewind_game(strong_parents, weak_parents, strong_to_move)
                        return game, weak_to_move
                    # Don't record stalemates
                else:
                    if (not new_weak_king_position.is_bishop_taken and
                            new_weak_king_position not in strong_parents):
                        strong_parents[new_weak_king_position] = strong_to_move
                        weak_parents[new_weak_king_position] = weak_to_move
                        next_round_unexplored.add(new_weak_king_position)
                    # Do nothing if weak king has taken bishop

        unexplored = next_round_unexplored


def main(board: Board):
    print("Looking for a checkmate")
    checkmate = search_for_checkmate(board)

    if checkmate is None:
        print("No checkmate found.")
        return

    game, checkmate_board = checkmate

    print("Found a checkmate")
    print("\nLegend:")
    print("    s   strong king")
    print("    w   weak king")
    print("    b   bishop")
    print("    |   white square")
    print("    -   black square\n")

    for side, board in zip(itertools.cycle(('Strong', 'Weak')), game):
        print(side, "to move:")
        print(board, end="\n\n")

    print("Checkmate!!")
    print(checkmate_board)


if __name__ == '__main__':
    main(Board(
        weak_king=Square(0, 4),
        strong_king=Square(7, 4),
        bishop1=Square(7, 2),
        bishop2=Square(7, 5)))

On this version, I purposefully added checks for bishop movements where the 3 other pieces could block the path:

1. it removes cases where a bishop could jump over the weak king, eliminating the need to check for impossible positions;
2. it allows two bishops on the same color (after pawn promotion, for instance);
3. it was as easy to check as the strong king alone.

Caching

As you can see, I dropped any kind of caching in the proposed version. But this is easy to put it back: a given square always generate the same positions for a king move or a bishop move. Let's cache the generated positions using functools.lru_cache.

But there is a catch, as the given code return a generator of positions, caching without caution will always return the same generator which will be exhausted after the first traversal. So we should adapt a bit and return sets as you did in you original code.

Last thing to consider is the size of the cache. We could be tempted to use 64 as there are only the square of the board that generate new positions or take into account invalid positions that may have been generated around the board to bump that limit to 100 squares. This is unnecessary as these Square will never call their **_movements methods, thus not poisoning the cache:

class Square(namedtuple('Square', 'row col')):
    def move(self, row, col):
        return Square(self.row + row, self.col + col)

    @property
    def is_valid(self):
        return all(x in BOARD_LIMITS for x in self)

    @lru_cache(maxsize=64)
    def king_movements(self):
        return set(self._king_movements())

    def _king_movements(self):
        for row, col in itertools.product(range(-1, 2), repeat=2):
            new_position = self.move(row, col)
            if new_position != self and new_position.is_valid:
                yield new_position

    @lru_cache(maxsize=64)
    def bishop_movements(self):
        return set(self._bishop_movements())

    def _bishop_movements(self):
        # First diagonal
        offset = min(self)
        yield from self._bishop_helper(-offset, 1)

        # Second diagonal
        offset = min(self.row, 7 - self.col)
        yield from self._bishop_helper(-offset, -1)

    def _bishop_helper(self, offset, column_adjustment):
        position = self.move(offset, column_adjustment * offset)
        while position.is_valid:
            if position != self:
                yield position
            position = position.move(1, column_adjustment)

Runtime just got down from 6m20s to 1m0s with no other modification of the code.