Python Amazon Checkmate

Question

Found this challenge on CodeFights and would love a review of my style.

Description

An amazon (also known as a queen+knight compound) is an imaginary chess piece that can move like a queen or a knight (or, equivalently, like a rook, bishop, or knight). The diagram below shows all squares which the amazon attacks from e4 (circles represent knight-like moves while crosses correspond to queen-like moves).

Recently you've come across a diagram with only three pieces left on the board: a white amazon, white king and black king. It's black's move. You don't have time to determine whether the game is over or not, but you'd like to figure it out in your head. Unfortunately, the diagram is smudged and you can't see the position of the black's king, so it looks like you'll have to check them all.

Given the positions of white pieces on a standard chessboard, determine the number of possible black king's positions such that:

• it's checkmate (i.e. black's king is under amazon's attack and it cannot make a valid move);
• it's check (i.e. black's king is under amazon's attack but it can reach a safe square in one move);
• it's stalemate (i.e. black's king is on a safe square but it cannot make a valid move);
• black's king is on a safe square and it can make a valid move.

Note that two kings cannot be placed on two adjacent squares (including two diagonally adjacent ones).

Example

For king = "d3" and amazon = "e4", the output should be amazonCheckmate(king, amazon) = [5, 21, 0, 29].

Code

def amazonCheckmate(king, amazon):

    # is in range of n for a and b
    def squared_range(a, b, n):
        return (a[0] >=  b[0] -n and a[0] <= b[0] + n and a[1] >= b[1]-n and a[1] <= b[1] + n)

    # if any square +1 // -1 for y and x are not under check ==> King escapes
    def no_valid_moves(a, b):
        if(a!=1):
            if([a-1,b] in notUnderCheck):
                return False
            if(b!=1):
                if([a-1,b-1] in notUnderCheck):
                    return False
            if(b!=8):
                if([a-1,b+1] in notUnderCheck):
                    return False        
        if(a!=8):
            if([a+1,b] in notUnderCheck):
                return False
            if(b!=1):
                if([a+1,b-1] in notUnderCheck):
                    return False
            if(b!=8):
                if([a+1,b+1] in notUnderCheck):
                    return False                
        if(b!=8):
            if([a,b+1] in notUnderCheck):
                return False
        if(b!=1):
            if([a,b-1] in notUnderCheck):
                return False        
        return True

    # declaring variables
    result = [0, 0, 0, 0]    
    letters = ['','a','b','c','d','e','f','g','h']
    notUnderCheck = []
    underCheck = []

    k_file = letters.index(king[0])
    k_col = int(king[1])
    a_file = letters.index(amazon[0])
    a_col = int(amazon[1])
    k = [k_file, k_col]
    q = [a_file, a_col]

    # if king defends queen the square where queens stand is undercheck else not under check
    if(squared_range(k, q, 1)):
        underCheck.append(q)
    else:
        notUnderCheck.append(q)

    # a grid 8x8 check which squares are under check and which are not
    for x in range(1, 9):
        for y in range(1, 9):

            # Squares to excldue are defended by King or postion of Amazon
            if(not squared_range([x,y], k, 1) and not [x,y] == q):                
                # if in deadly square of queen 5x5
                # 
                if(squared_range([x,y],q, 2)):
                    underCheck.append([x,y])

                # Check if on the same file and not if king is in between                
                elif (x == a_file):
                    if(not (k_file == a_file and (y > k_col > a_col or y < k_col < a_col))):
                        underCheck.append([x,y])

                # Check if on the same column and not if king in between
                elif( y == a_col):
                    if(not (k_col == a_col and ( x > k_file > a_file or x < k_file < a_file))):
                        underCheck.append([x,y])

                # Check diagonal and not if king in between
                # Black_King on Diagonaal van Queen
                elif(abs(x - a_file) == abs(y - a_col)):
                     if( not(abs(k_file - a_file) == abs(k_col - a_col) 
                         and ( ( x < k_file < a_file and y < k_col < a_col) 
                             or (x < k_file < a_file and y > k_col > a_col)
                             or (x > k_file > a_file and y < k_col < a_col)
                             or (x > k_file > a_file and y < k_col < a_col) ) ) ):
                        underCheck.append([x,y])
                else:
                    notUnderCheck.append([x, y])

            # Add the squares where to White_King stands to checksquares
            elif(squared_range([x,y], k, 1)):
                underCheck.append([x,y])

    # for each cell in grid check surounding cells strengths update result accordingly
    for x in range(1, 9):
        for y in range(1, 9):
            # Exclude q and kings range
            if(not squared_range([x,y], k, 1) and not [x,y] == q):
                # if current square under Check
                if([x, y] in underCheck):
                    # if no possible moves result[0]++
                    # else result[1]++
                    if(no_valid_moves(x, y)):
                        print("Checkmate at: [" + str(x) + ", " + str(y) + "]")
                        result[0] += 1
                    else:
                        print("Check at: [" + str(x) + ", " + str(y) + "]")
                        result[1] += 1                                                    
                else:
                    # if no possible moves result[2]++
                    # else result[3]++
                    if(no_valid_moves(x, y)):
                        print("Stuck at: [" + str(x) + ", " + str(y) + "]")
                        result[2] += 1
                    else:
                        print("Safe at: [" + str(x) + ", " + str(y) + "]")
                        result[3] += 1     
    return result

Question

1. no_valid_moves looks clunky as hell! How can this be made beautifull again?
2. Is my general logic ok, or should I have divided the problem in easier to comprehend functions?
3. Any stylistic review is welcome too.

Answer 1

Clunky methods can be the result of awkward setups. In your case, yes, you should have divided the problem up in smaller chunks. Your algorithm is so tightly bound to the current problem that

1. it'd be extremely hard to adapt it later on, say, should you need to add another piece in the equation, and
2. it's likely your logic has bugs (it does) because it's tough to cover all corner cases when you have to explicitly check all cases, instead of allowing them to emerge as a consequence of combining the solution to smaller problems.

Your problem can be broken down as:

1. find all the squares under attack by the white pieces
2. generate all valid black king squares
3. for each square make two checks: is the black king under attack and does he have any safe squares to run to?

1. Can be further broken down to generating the attacking squares for each white piece individually and taking their union at the end. For example, generating the squares a king can move into contributes to all 3 steps above. For step 2 it is by figuring out which square placements are not valid for the black king, since he can't start at any squares the white king can move into.

In terms of movements, there are two types of pieces we are concerned with. Ray-like pieces, i.e., bishops and rooks, can move outwards in various directions, but their movement can be blocked if another piece is in their path. And we also have occupying pieces, i.e., king and knight, which can influence certain squares regardless of any piece arrangement.

To generate the amazon moves, I would write rules for a bishop, rook and knight individually. Then, the first two would have to be filtered depending on whether the white king is in the amazon's way.

---

In terms of stylistic comments

• write your functions/methods in snake_case. CamelCase should be reserved for classes only.
• Don't exceed 80 characters per line.
• Use docstrings to describe what a function does or how to use it. Comments aren't meant for this.
• Avoid using comments to describe what the code does when it's trivial. The code should ideally explain itself and comments should be reserved for why you did something the way you did. Or if you're using a complex algorithm, a summary of its implementation.
• Don't use brackets in statements when it isn't required, e.g., if a in b: is fine, while if (a in b): is visual clutter.
• Use descriptive names. If the code is complicated, it's hard to keep track what k_col and k represent. Also, the chess terms for rows/columns are rank/file. So, I would have used king_file, king_rank and king_coords.
• Use str.format() for complicated strings instead of concatenating them together. For example, 'Checkmate at [{}, {}]'.format(x, y).
• A function that returns a boolean should have a name like is_something, can_do_something, has_something, etc.
• You hardcode magic numbers, e.g, the size of the board. It's better if you define a constant (all-caps) at the top, so if you ever need to change it, it'll be trivial and error-free.
• You convert the files a-h to the numbers 1-9, when it's more natural to work with 0-based indices for iterables. letters = ['','a','b','c','d','e','f','g','h'] is an obvious hack, trying to fight against what is most convenient.
• You should consider validating the inputs, because nothing prevents someone from feeding in 'j9', or other garbage. Or that the king's and amazon's positions don't overlap.

---

Overall, here is how I would have tackled it.

import itertools
from string import ascii_lowercase as files

SIZE = 8

squares = tuple(''.join(square[::-1]) for square in
               itertools.product(map(str, range(1, SIZE+1)), files[:SIZE]))
coords = tuple(itertools.product(range(SIZE), range(SIZE)))

def is_in_bounds(square):
    return 0 <= square < SIZE

def king_moves(rank, file):
    moves = []
    if rank - 1 >= 0:
        moves += [(rank-1, f) for f in range(file-1, file+2) if is_in_bounds(f)]
    moves += [(rank, f) for f in (file-1, file+1) if is_in_bounds(f)]
    if rank + 1 < SIZE:
        moves += [(rank+1, f) for f in range(file-1, file+2) if is_in_bounds(f)]
    return moves

def rook_moves(rank, file):
    moves = [[(r, file) for r in range(rank-1, -1, -1)]]
    moves += [[(r, file) for r in range(rank+1, SIZE)]]
    moves += [[(rank, f) for f in range(file-1, -1, -1)]]
    moves += [[(rank, f) for f in range(file+1, SIZE)]]
    return moves

def bishop_moves(rank, file):
    down = range(-1, -rank-1, -1)
    up = range(1, SIZE-rank)
    moves = [[(rank+i, file-i) for i in down if is_in_bounds(file-i)]]
    moves += [[(rank+i, file+i) for i in down if is_in_bounds(file+i)]]
    moves += [[(rank+i, file-i) for i in up if is_in_bounds(file-i)]]
    moves += [[(rank+i, file+i) for i in up if is_in_bounds(file+i)]]
    return moves

def knight_moves(rank, file):
    offsets = ((-2, -1), (-2, 1),
               (-1, -2), (-1, 2),
               (1, -2), (1, 2),
               (2, -1), (2, 1))
    moves = [(rank+x, file+y) for x, y in offsets
             if is_in_bounds(rank+x) and is_in_bounds(file+y)]
    return moves

def filter_ray_attacks(moves, piece_coords):
    filtered_moves = []
    for direction in moves:
        if piece_coords in direction:
            # +1 because it can influence the friendly occupied square
            direction = direction[:direction.index(piece_coords)+1]
        filtered_moves += direction
    return filtered_moves

def amazon_moves(rank, file, king):
    moves = (filter_ray_attacks(rook_moves(rank, file), king) +
             filter_ray_attacks(bishop_moves(rank, file), king) +
             knight_moves(rank, file))
    return moves

def has_escape_squares(square, piece_moves, attacks):
    for move in piece_moves:
        if move not in attacks:
            return True
    return False

def solve_amazon_problem(king, amazon):
    for piece in (king, amazon):
        if piece not in squares:
            raise ValueError('{} is not a valid square.'.format(piece))
    if king == amazon:
        raise ValueError('The king and amazon must not occupy the same square.')

    king = coords[squares.index(king)]
    amazon = coords[squares.index(amazon)]

    king_attacks = king_moves(*king)
    all_attacks = set(king_attacks + amazon_moves(*amazon, king))

    black_not_allowed_squares = set([king, amazon] + king_attacks)
    black_king_positions = (square for square in coords
                            if square not in black_not_allowed_squares)

    # `result` is 4-item list:
    # 0 - checkmate
    # 1 - check
    # 2 - stalemate
    # 3 - safe
    # 0-1 means the king is initially in check and 2-3 not.
    # Therefore, we define `check_offset` as either 0 or 2.
    # Between 0 and 1 (or 2 and 3) the difference is whether the king
    # has any safe moves. This is represented by the `escape_offset`, which
    # is either 0 or 1.
    # The final state is then `check_offset + escape_offset`.
    result = [0] * 4
    for square in black_king_positions:
        black_king_moves = king_moves(*square)
        check_offset = 2 * (1 - (square in all_attacks))
        escape_offset = has_escape_squares(
            square, black_king_moves, all_attacks)
        result[check_offset + escape_offset] += 1
    return result

By the way, for various inputs your code seems to produce incorrect results. For example, for king at 'f7' and amazon at 'e7', your program claims there is a stalemate when the king starts at 'h8', while 'h7' is a safe escape. It's likely you didn't properly calculate that the king blocks the amazon's movements along the 7th rank right from the king.

Answer 2

I think the original proposal as well as answers given so far are not in adequacy with the simplicity, and even statement, of the problem. As it is formulated, the most natural answer is just to loop over the 8 x 8 squares and increment a counter for the respective case, for each allowed square:

def amazonCheckmate(king: str, amazon: str, n_rows = 8, n_cols = None):
    """For given positions of the white king and amazon (moves like queen or
       knight), return [c0,c1,c2,c3] where c0-c3 count the number of squares
       on which the black king is checkmate, in check, stalemate or safe,
       respectively, when the white king and amazon are at the positions
       given by the arguments in usual notation, e.g., 'a1' or 'e4', etc.
       Number of columns (files) defaults to number of rows of the board."""
    # Convert coordinate string to (col, row) (= file, rank): Consider it
    # as a base 20 number (to allow 'a', 'b',... as digits 10, 11,...);
    # subtract 10*20 + 1 so that 'a' and '1' correspond to col. and row = 0:
    if n_cols == None:
       n_cols = n_rows
    kc, kr = divmod(int( king , 20)-201, 20) 
    ac, ar = divmod(int(amazon, 20)-201, 20)
    cnt = [0]*4   # counter for each of: checkmate, check, stalemate, safe.
    for r in range(n_rows):      # row or "rank"
        for c in range(n_cols):  # column or "file"
            if not forbidden(r, c):  # i.e.: too close to enemy king
                cnt[ (0 if threatened(r, c) else 2) + can_move(r, c) ] += 1
    return cnt

Now we just need the simple if not trivial functions forbidden() (too close to the enemy king), threatened() and can_move(). If we insert them somewhere before the above loop, they can access the coordinates (kr,kc) and (ar,rc) of the white king and amazon as variables defined within their environment:

    def forbidden(r, c):
        """True if (r,c) is at distance <= 1 from white king."""
        return abs(r-kr) < 2 and abs(c-kc) < 2

    def can_move(r, c):
        """Is there a safe square next to (r,c)?"""
        return any(not threatened(rr,cc) and not forbidden(rr,cc)
                   for rr in range(max(r-1,0), min(r+2, n_rows))
                   for cc in range(max(c-1,0), min(c+2, n_cols),
                                   2 if rr == r else 1))  # skip rr,cc = r,c

    def threatened(r, c):
        """Is square (r,c) "threatened"/controlled by the amazon (or king)?
           That's the case if (r,c) is at distance <= 2 from the amazon, 
           or on the same row or column or diagonal as the amazon,
           and the white king is not in between the two."""
        return (abs(r-ar) < 3 and abs(c-ac) < 3) or (r==ar or c==ac or 
                abs(c-ac)==abs(r-ar)) and not shielded(r,c)

    def shielded(r, c):
        """Is the enemy king between (r,c) and (ar,ac), in diagonal or
        straight line? Equivalently: the directional vectors
        (delta row, delta col) collinear <=> determinant is zero."""
        return (kc-c)*(ar-r) == (kr-r)*(ac-c) and \
               min(c,ac) <= kc <= max(c,ac) and min(r,ar) <= kr <= max(r,ar)

That's all that is needed to solve the problem in a simple, self-explaining way.