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

    gears = dict()

    for y in range(1, len(lines) - 1):
        n = 0
        ni = []
        for x, c in enumerate(lines[y]):
            if c in "0123456789":
                n = n * 10 + int(c)
                ni.append(x)
            elif n != 0:
                for i in range(min(ni) - 1, max(ni) + 2):
                    if lines[y-1][i] == "*":
                        if (i, y-1) not in gears:
                            gears[(i, y-1)] = []
                        gears[(i, y-1)].append(n)
                    if lines[y][i] == "*":
                        if (i, y) not in gears:
                            gears[(i, y)] = []
                        gears[(i, y)].append(n)
                    if lines[y+1][i] == "*":
                        if (i, y+1) not in gears:
                            gears[(i, y+1)] = []
                        gears[(i, y+1)].append(n)

                n = 0
                ni = []

    return sum([ns[0] * ns[1] if len(ns) == 2 else 0 for ns in gears.values()])


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

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