aoc/2023/03/solution.py

#!/usr/bin/python3
import sys
import re

def generate_field(span, row, len_col, len_row):
    left = max(span[0] - 1, 0)
    right = min(span[1], len_row - 1)
    up = max(row - 1, 0)
    down = min(row + 1, len_col - 1)
    field = [
        (up, left), (up, right), (down, left), (down, right),
    ]
    return field

def check_field(grid, field):
    for y in range(field[0][0], field[2][0] + 1):
        for x in range(field[0][1], field[1][1] + 1):
            if grid[y][x] not in '.0123456789':
                return True
    return False

def find_gear(grid, field):
    for y in range(field[0][0], field[2][0] + 1):
        for x in range(field[0][1], field[1][1] + 1):
            if grid[y][x] in '*':
                return (y, x)
    return (-1, -1)

def part1():
    print('part 1')
    grid = []
    numbers = []
    spans = []
    rows = []
    for row, line in enumerate(sys.stdin):
        grid.append(list(line.strip()))
        for match in re.finditer(r'\d+', line.strip()):
            numbers.append(int(match.group()))
            spans.append(match.span())
            rows.append(row)
    summed = 0
    for number, row, span in zip(numbers, rows, spans):
        field = generate_field(span, row, len(grid), len(grid[row]))
        #print(number, row, span, field)
        if check_field(grid, field):
            summed += number
        #else:
            #print(f'{number} not in field')
    print(summed)

def part2():
    print('part 2')
    grid = []
    numbers = []
    spans = []
    rows = []

    for row, line in enumerate(sys.stdin):
        grid.append(list(line.strip()))
        for match in re.finditer(r'\d+', line.strip()):
            numbers.append(int(match.group()))
            spans.append(match.span())
            rows.append(row)

    gears = {}
    for number, row, span in zip(numbers, rows, spans):
        field = generate_field(span, row, len(grid), len(grid[row]))
        y, x = find_gear(grid, field)
        if y > 0 and x > 0:
            if (y, x) in gears:
                gears[(y, x)].append(number)
            else:
                gears[(y, x)] = [number]
    
    ratios = 0
    for gear in gears:
        if len(gears[gear]) > 1:
            ratio = 1
            for number in gears[gear]:
                ratio *= number
            ratios += ratio
            
    print(gears)
    print(ratios)

if sys.argv[1] in '1':
    part1()
else:
    part2()