Files
advent-of-code-2023/advent_of_code/gears.py
2023-12-16 17:40:14 -08:00

123 lines
3.2 KiB
Python

"Day 3: Gear Ratios"
from __future__ import annotations
from collections import UserDict
from dataclasses import dataclass
from functools import reduce
import math
from typing import NamedTuple
class Symbol(NamedTuple):
value: str
row: int
col: int
class Number(NamedTuple):
value: str
row: int
col: int
anchor: Symbol | None = None
def extend_digit(self, digit: str) -> Number:
return self._replace(value=self.value + digit)
class Part(NamedTuple):
number: Number
symbol: Symbol
@dataclass
class Schematic:
numbers: list[Number]
symbols: list[Symbol]
@classmethod
def parse(cls, input: str) -> Schematic:
row = 0
col = 0
current_number: Number | None = None
numbers: list[Number] = []
symbols: list[Symbol] = []
for char in input:
match char:
# Digit
case n if n in "0123456789":
if not current_number:
current_number = Number("", row, col)
current_number = current_number.extend_digit(char)
col += 1
# Blank
case ".":
if current_number:
numbers.append(current_number)
current_number = None
col += 1
# Newline
case "\n":
if current_number:
numbers.append(current_number)
current_number = None
row += 1
col = 0
# Schematic smybol
case _:
if current_number:
numbers.append(current_number)
current_number = None
symbols.append(Symbol(char, row, col))
col += 1
# Finalize a number at the end of the schematic
if current_number:
numbers.append(current_number)
return cls(numbers, symbols)
def parts(self) -> set[Part]:
return {
Part(number, symbol)
for number in self.numbers
for symbol in self.symbols
if (
# Symbol within 1 row
(number.row - 1 <= symbol.row <= number.row + 1)
and
# Symbol within 1 column
(number.col - 1 <= symbol.col <= number.col + len(number.value))
)
}
def part_groups(self) -> dict[Symbol, set[Part]]:
groups: dict[Symbol, set[Part]] = {}
for part in self.parts():
if part.symbol not in groups:
groups[part.symbol] = set()
groups[part.symbol].add(part)
return groups
def solve_part_1(input: str) -> int:
schematic = Schematic.parse(input)
return sum(int(part.number.value) for part in schematic.parts())
def solve_part_2(input: str) -> int:
schematic = Schematic.parse(input)
return sum(
reduce(
lambda x, y: x * y,
(int(part.number.value) for part in parts),
)
for symbol, parts in schematic.part_groups().items()
if symbol.value == "*"
if len(parts) == 2
)