AOC 23 - Day 3

main.py

import collections
import copy
import dataclasses
import enum
import re
import pprint
import string
import sys


GEAR_COUNT = 1


class SymbolType(enum.Enum):
    NUMBER = 1
    SYMBOL = 2
    BLANK = 3


@dataclasses.dataclass
class Symbol:
    c: str
    t: SymbolType
    m: int | None


def symtype(c: str) -> SymbolType:
    if c in "0123456789":
        return SymbolType.NUMBER
    elif c == ".":
        return SymbolType.BLANK
    else:
        return SymbolType.SYMBOL


def make_symbol(c: str) -> Symbol:
    return Symbol(c=c, t=symtype(c), m=None)


def neighbors(
    i: int, j: int, i_range: int, j_range: int, diagonals: bool = True
) -> list[tuple[int]]:
    ns = []
    if i > 0 and j > 0:
        ns.append((i - 1, j - 1))
    if i > 0:
        ns.append((i - 1, j))
    if i > 0 and j < j_range - 1:
        ns.append((i - 1, j + 1))
    if j > 0:
        ns.append((i, j - 1))
    if j < j_range - 1:
        ns.append((i, j + 1))
    if i < i_range - 1 and j > 0:
        ns.append((i + 1, j - 1))
    if i < i_range - 1:
        ns.append((i + 1, j))
    if i < i_range - 1 and j < j_range - 1:
        ns.append((i + 1, j + 1))
    return ns


def mark_numbers(syms: list[list[Symbol]]) -> list[list[Symbol]]:
    global GEAR_COUNT

    i_range = len(syms)
    j_range = len(syms[0])
    out = copy.deepcopy(syms)

    for i in range(i_range):
        for j in range(j_range):
            if out[i][j].t == SymbolType.SYMBOL:
                if out[i][j].m is None:
                    if out[i][j].c == "*":
                        out[i][j].m = GEAR_COUNT
                        GEAR_COUNT += 1
                    else:
                        out[i][j].m = 0
                ns = neighbors(i, j, i_range, j_range)
                for i_, j_ in ns:
                    if out[i_][j_].t == SymbolType.NUMBER:
                        out[i_][j_].m = out[i][j].m
            elif out[i][j].t == SymbolType.NUMBER and out[i][j].m is not None:
                ns = neighbors(i, j, i_range, j_range)
                for i_, j_ in ns:
                    if out[i_][j_].t == SymbolType.NUMBER:
                        out[i_][j_].m = out[i][j].m

    return out


def collect_marked(syms: list[list[Symbol]]) -> list[int]:
    marked = []

    for line in syms:
        acc = ""
        for sym in line:
            if sym.t == SymbolType.NUMBER and sym.m is not None:
                acc += sym.c
            elif acc:
                marked.append(int(acc))
                acc = ""
        if acc:
            marked.append(int(acc))
            acc = ""
    return marked


def collect_gears(syms: list[list[Symbol]]) -> list[tuple[int]]:
    gears = []

    for line in syms:
        acc = ""
        gear = None
        for sym in line:
            if sym.t == SymbolType.NUMBER and sym.m is not None:
                acc += sym.c
                gear = sym.m
            elif acc:
                if gear != 0:
                    gears.append((gear, int(acc)))
                acc = ""
                gear = None
        if acc:
            if gear != 0:
                gears.append((gear, int(acc)))
            acc = ""
            gear = None
    return gears


def main():
    lines = [line.strip() for line in open(sys.argv[1]).readlines()]
    syms = [[make_symbol(c) for c in line] for line in lines]

    prev = syms
    marks = mark_numbers(prev)
    while prev != marks:
        prev = marks
        marks = mark_numbers(prev)

    gears = collections.defaultdict(list)
    for g, num in collect_gears(marks):
        gears[g].append(num)

    print(sum(collect_marked(marks)))
    print(sum([v[0] * v[1] for v in gears.values() if len(v) == 2]))


if __name__ == "__main__":
    main()