p03

#!/usr/bin/env python3
import math
from dataclasses import dataclass
import fileinput

@dataclass
class NUMBER:
    value: int
    length: int
    coords: tuple[int, int]

    def is_adjacent(self, coords) -> bool:
        y, sx = self.coords
        for x in range(sx, sx + self.length):
            if (y, x + 1) in coords or \
               (y, x - 1) in coords or \
               (y + 1, x) in coords or \
               (y + 1, x + 1) in coords or \
               (y + 1, x - 1) in coords or \
               (y - 1, x) in coords or \
               (y - 1, x + 1) in coords or \
               (y - 1, x - 1) in coords:
                return True

        return False

symbols = {}
numbers = []
number = []

def addnum(y: int, x: int) -> None:
    if number:
        ln = len(number)
        numbers.append(NUMBER(int(''.join(number)), ln, (y, x - ln)))
        number.clear()

for y, line in enumerate(fileinput.input()):
    line = line.strip()
    for x, c in enumerate(line):
        if c == '.':
            addnum(y, x)
        elif c.isdigit():
            number.append(c)
        else:
            addnum(y, x)
            symbols[(y, x)] = c

    addnum(y, len(line))

print('P1 =', sum(n.value for n in numbers if n.is_adjacent(symbols)))

tot = 0
for coords, symbol in symbols.items():
    if symbol == '*':
        nums = [n.value for n in numbers if n.is_adjacent([coords])]
        if len(nums) == 2:
            tot += math.prod(nums)

print('P2 =', tot)