StyleInCode

""" day_04_01.py """

# usage: python3 day_04_01.py < input

import sys


def locate(letter, grid):
    """ coordinate of each letter """
    return [(i, j) for j, row in enumerate(grid)
            for i, col in enumerate(row) if col == letter]


def words(pos, num, grid):
    """ num letter words - in all directions """
    i, j = pos

    directions = [(1, 0), (-1, 0), (0, 1), (0, -1),
                  (1, 1), (-1, -1), (1, -1), (-1, 1)]

    output = []
    for dx, dy in directions:
        if i + num * dx < -1 or j + num * dy < -1:
            continue
        if i + num * dx > len(grid[0]) or j + num * dy > len(grid):
            continue
        output.append(''.join([grid[j + d * dy][i + d * dx]
                               for d in range(num)]))
    return output


with sys.stdin as infile:
    puzzle = [list(line.strip()) for line in infile]

investigate = locate('X', puzzle)

total = 0

for location in investigate:
    total += words(location, 4, puzzle).count('XMAS')

print(total)

""" day_04_02.py """

# usage: python3 day_04_02.py < input

import sys


def locate(letter, grid):
    """ coordinate of each letter """
    return [(i, j) for j, row in enumerate(grid)
            for i, col in enumerate(row) if col == letter]


def words(pos, grid):
    """ 3 letter words - in a cross """
    i, j = pos

    directions = [(1, 1), (1, -1)]

    output = []
    for dx, dy in directions:
        if i - 2 < -1 or j - 2 < -1:
            continue
        if i + 2 > len(grid[0]) or j + 2 > len(grid):
            continue
        output.append(''.join([grid[j + d * dy][i + d * dx]
                               for d in range(-1, 2)]))
    return output


with sys.stdin as infile:
    puzzle = [list(line.strip()) for line in infile]

investigate = locate('A', puzzle)

total = 0

for location in investigate:
    cross = words(location, puzzle)
    if cross:
        for word in cross:
            if word not in ['MAS', 'SAM']:
                break
        else:
            total += 1

print(total)