import re


def do_task(lines):
    lines = (
        ["." * (len(lines[0]) + 2)]
        + ["." + line.strip() + "." for line in lines]
        + ["." * (len(lines[0]) + 2)]
    )

    symbol_sum = 0

    for y in range(1, len(lines) - 1):
        for m in re.finditer(r"\d+", lines[y]):
            if re.search(
                r"[^.0-9]",
                lines[y - 1][m.start() - 1 : m.end() + 1]
                + lines[y][m.start() - 1] + lines[y][m.end()]
                + lines[y + 1][m.start() - 1 : m.end() + 1],
            ):
                symbol_sum += int(m[0])

    return symbol_sum


with open("day03/data/input100.txt") as infile:
    lines = infile.readlines()
    result = do_task(lines)
    print(result)
    print(result == 4361)

with open("day03/data/input101.txt") as infile:
    lines = infile.readlines()
    result = do_task(lines)
    print(result)
    print(result == 525181)