day03.py

from itertools import product
from math import prod

numbers = []
grid = open("input/day03.txt").read().split("\n")

def is_symbol(x, y):
    c = grid[max(0, min(len(grid)-1, x))][max(0, min(len(grid[0])-1, y))]
    return not c.isdigit() and c != "."

for x, line in enumerate(grid):
    y = 0
    while y < len(line):
        j = y
        if line[j].isdigit():
            is_part = is_symbol(x-1, j-1) or is_symbol(x+1, j-1) or is_symbol(x, j-1)
            while j < len(line) and line[j].isdigit():
                is_part = is_part or is_symbol(x-1, j) or is_symbol(x+1, j)
                j += 1
            if y != j:
                is_part = is_part or is_symbol(x-1, j) or is_symbol(x+1, j) or is_symbol(x, j)
                if is_part:
                    numbers.append((x, y, line[y:j]))
        y = j + 1

print(f"a = {sum(int(n[2]) for n in numbers)}")

# part 2
def get_number(x, y):
    y = max(0, min(len(grid[0])-1, y))
    line = grid[max(0, min(len(grid)-1, x))]
    c = line[y]
    if not c.isdigit():
        return None
    begin = y
    while begin >= 0 and line[begin].isdigit():
        begin -= 1
    if not line[begin].isdigit():
        begin += 1
    end = y
    while end < len(line) and line[end].isdigit():
        end += 1
    return begin, end, line[begin:end]

res = 0
for x, line in enumerate(grid):
    for y, c in enumerate(line):
        if c == "*":
            numbers = set()
            for dx, dy in product([-1, 0, 1], [-1, 0, 1]):
                number = get_number(x+dx, y+dy)
                if number is not None:
                    numbers.add(number)
            if len(numbers) == 2:
                res += prod(int(number[2]) for number in numbers)
print(f"b = {res}")