adventofcode2023/day03/sol.py

import sys

input = sys.stdin
data = [line.strip() for line in input]

def issymbol(c):
    return c != '.' and not c.isalnum()

def nexttosymbol(row, col):
    for x in [col-1, col, col+1]:
        for y in [row-1, row, row+1]:
            if 0 <= y < len(data) and 0 <= x < len(data[y]):
                if issymbol(data[y][x]):
                    return True
    return False

def findgears(row, col):
    for x in [col-1, col, col+1]:
        for y in [row-1, row, row+1]:
            if 0 <= y < len(data) and 0 <= x < len(data[y]):
                if data[y][x] == '*':
                    yield y,x


part1 = 0
gears = {}
for no, line in enumerate(data):
    places = []
    digits = []
    i = 0
    while i < len(line):
        if line[i].isdigit() and (i == 0 or not line[i-1].isdigit()):
            j = i+1
            while j < len(line) and line[j].isdigit():
                j += 1
            digits.append(int(line[i:j]))
            places.append(list(range(i,j)))
            i = j
        else:
            i += 1

    for n, pos in zip(digits, places):
        if any(nexttosymbol(no, i) for i in pos):
            print(n)
            part1 += n

        nearby_gears = set()
        for i in pos:
            nearby_gears.update(findgears(no, i))
        for g in nearby_gears:
            gears.setdefault(g, []).append(n)

print(part1)

part2 = 0
for g, nums in gears.items():
    if len(nums) == 2:
        part2 += nums[0] * nums[1]
print(part2)