Find the number of pawns available for a rook to capture

Question

I find that very very often with matrix questions, the traversal and comparison part of the solution is very verbose and repetitive.

For example, this code (written for a coding challenge) to find the number of pawns available for a rook to capture in one move on a chess board (R is a white rook, B is a white bishop, and p is a black pawn):

class Solution:
    def numRookCaptures(self, board: List[List[str]]) -> int:
        for r in range(len(board)):
            for c in range(len(board[r])):
                if board[r][c] == 'R':
                    return self.find_p(board, r, c)

    def find_p(self, board, r, c):
        count = 0
        for i in range(r + 1, 8):
            if board[i][c] in ('B', 'p'):
                count += int(board[i][c] == 'p')
                break
        for i in range(r - 1, -1, -1):
            if board[i][c] in ('B', 'p'):
                count += int(board[i][c] == 'p')
                break
        for i in range(c + 1, 8):
            if board[r][i] in ('B', 'p'):
                count += int(board[r][i] == 'p')
                break
        for i in range(c - 1, -1, -1):
            if board[r][i] in ('B', 'p'):
                count += int(board[r][i] == 'p')
                break
        return count

This is the most concise way of writing it that I've found, without compromising clarity. The in for comparison and int for incrementing the count each save us a few lines of code.

But, it's still really verbose and repetitive. Each of the for loops are essentially doing the same exact thing, just in a different direction.

This comes up in every matrix question, whether it's just linear checks like above, BFS, DFS, they pretty much all have a four-directional comparison and traversal that produces 4 copies of code that are all identical in function.

Is there a less verbose way of writing these, without being overly fancy and hard-to-read?

Answer 1

Dimensions

How big is your chess board, and is it square, or at least rectangular?

First, you use range(len(board)) to determine the rows of the chess board, but then you use range(len(board[r])) for the columns, which means you can handle each row having a different number of columns!

(Supporting a non-rectangular chess board would be a challenging challenge, but totally possible, as demonstrated above.)

But then, you use range(r+1, 8) and range(c+1, 8), indicating that you only support square 8x8 chess boards.

Unfortunately, you haven’t posted the challenge problem text, so we can only guess at the actual challenge. Additional context would go a long way.

Verbosity & Fancyness

In order to be less verbose, we need to do something at least a little bit “fancy”. With only rook movements, it may seem overly fancy, but if you later support bishops and queens, the additional capability we build in will be leveraged, and very little code need be added.

Additionally, you only need the white bishop blocking the white rook, and only black pawns being captured. What about white pawns with a black rook and bishop? This shouldn’t be a special case; it should be generalized. But again, generalization one one hand will be more verbose, and “fancy”, and the other hand will be leveraged to make the code agnostic to the current player.

First of all, you should have an on_board function, which will return True if a given position is actually on the chess board. I’ll hard-code the board size, and assume the board is stored in self; the dimensions could be stored in self or derived from self.board on the fly.

def on_board(self, r, c):
    return r in range(8) and c in range(8)

Next, pieces get “blocked” by other pieces. So let’s formalize that. Again, I’ll assume the piece is the 'R', so any white (uppercase) piece will block it. (You should return the appropriate expression based on the colour of piece ... left to student)

def blocked(self, piece, r, c):
    return self.board[r][c].isupper()

A black pawn doesn’t “block” a white rook’s move, but it does terminate it in a “capture” event. Again, we’ll assume a white piece is moving, so it will “capture” any black (lowercase) piece. (Again, base actual return value on piece ... left to student)

def capture(self, piece, r, c):
    return self.board[r][c].islower()

Rooks, bishops, and queens can move in a straight line any number of spaces until blocked, or it captures another piece. Note that a blocked move is not a valid move, but a capture move is a valid move.

def direction_moves(self, piece, r, c, dr, dc):
    r += dr
    c += dc
    while self.on_board(r, c) and not self.blocked(piece, r, c):
        yield r, c
        if self.capture(piece, r, c):
            break
        r += dr
        c += dc

Rooks, bishops and queens can move in more than one direction, so let’s allow multiple directions to be given at once:

def linear_moves(self, piece, r, c, *dirs):
    for dr, dc in dirs:
        yield from self.direction_moves(piece, r, c, dr, dc)

Now that we have the building blocks, getting a rook’s moves from a given position is easy:

def rook_moves(self, piece, r, c):
    yield from self.linear_moves(piece, r, c, (1, 0), (0, 1), (-1, 0), (0, -1))

Moves for the piece at a given position might be:

def moves(self, r, c):
    piece = self.board[r][c]
    if piece in {'R', 'r'}:   # Black or white rook
        yield from self.rook_moves(piece, r, c)
    # ... other types of pieces ...

With this, your find_p becomes trivial:

def find_p(self, r, c):
    return sum(1 for i, j in self.moves(r, c) if self.board[i][j] == 'p')

---

The above has increased the amount of code, but at the same time it is more generalized. It has work to be done to generalize it for white -vs- black. It would be trivial to add movements for bishops and queens to this structure, and that would not be repetitive code, unlike the original structure.

Unfortunately, you may conclude that it has become overly “fancy”.