chess_positions.py

# Consider a rectangular chess board of a × b squares. 
# For each of the squares imagine a lone queen standing on that square and 
# compute the number of squares under the queen's attack. Add all the numbers 
# you got for each of the a × b possible queen's positions.

# Example

# For a = 2 and b = 3, the output should be
# chessBoardSquaresUnderQueenAttack(a, b) = 26.

# Here are the squares that are under Queen attack:



# Input/Output

# [execution time limit] 4 seconds (py)

# [input] integer a

# Positive integer.

# Guaranteed constraints:
# 1 ≤ a ≤ 20.

# [input] integer b

# Positive integer.

# Guaranteed constraints:
# 1 ≤ b ≤ 20.

# [output] integer

from collections import namedtuple

Position = namedtuple('Position', 'x y')


def get_positions(a, b):
    positions = []
    for a in range(1, a + 1):
        for b in range(1, b + 1):
            positions.append(Position(a, b))
    return positions


def position_diag_1(pos):
    return Position(pos.x - 1, pos.y - 1)


def position_diag_2(pos):
    return Position(pos.x + 1, pos.y + 1)


def position_diag_3(pos):
    return Position(pos.x + 1, pos.y - 1)


def position_diag_4(pos):
    return Position(pos.x - 1, pos.y + 1)


def calculate_diagonal(positions, queen_pos, fnc):
    p = queen_pos
    contagem = 0
    while (p in positions):
        p = fnc(p)
        if p in positions:
            print(p)
            contagem += 1
    return contagem


def get_diagonal(a, b, pos):
    positions = get_positions(a, b)
    contagem = calculate_diagonal(positions, pos, position_diag_1)
    contagem += calculate_diagonal(positions, pos, position_diag_2)
    contagem += calculate_diagonal(positions, pos, position_diag_3)
    contagem += calculate_diagonal(positions, pos, position_diag_4)
    return contagem


def chessBoardSquaresUnderQueenAttack(a, b):
    squares_under_attack = []
    positions = get_positions(a, b)
    
    for po in positions:
        # horizontal
        squares_under_attack.append(a - 1)
        # vertical
        squares_under_attack.append(b - 1)
        # diagonal
        squares_under_attack.append(get_diagonal(a, b, po))
        
    return sum(squares_under_attack)