def check_direction(lines: list[str], row_idx: int, col_idx: int, row_step: int, column_step: int) -> bool: column_count = len(lines[0]) row_count = len(lines) for step, letter in enumerate("MAS"): current_row = row_idx + (step + 1) * row_step current_col = col_idx + (step + 1) * column_step if current_row >= row_count or current_row < 0: return False if current_col >= column_count or current_col < 0: return False if lines[current_row][current_col] != letter: return False return True def get_current_pos_words(lines: list[str], row_idx: int, col_idx: int) -> int: if lines[row_idx][col_idx] != 'X': return 0 count = 0 for row_step in range(-1, 2): for column_step in range(-1, 2): if row_step == 0 and column_step == 0: continue if check_direction(lines, row_idx, col_idx, row_step, column_step): count += 1 return count data = [] with open("input") as f: for line in f: data.append(line.strip()) result = 0 for row, line in enumerate(data): for col in range(len(line)): result += get_current_pos_words(data, row, col) print(result)