aoc2023/aoc/day11.py

141 lines
4 KiB
Python

from __future__ import annotations
from array import array
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from itertools import accumulate, combinations
from .aoc import Aoc2, AocSameParser
@dataclass
class Position:
x: int
y: int
def __add__(self, other: Position) -> Position:
return Position(self.x + other.x, self.y + other.y)
def __sub__(self, other: Position) -> Position:
return Position(self.x - other.x, self.y - other.y)
def __or__(self, other: Position) -> int:
return abs(self.x - other.x) + abs(self.y - other.y)
def __rmul__(self, other: int) -> Position:
return Position(other * self.x, other * self.y)
def __getitem__(self, pi: int) -> int:
if pi == 0:
return self.x
return self.y
def __repr__(self) -> str:
return f"[{self.x},{self.y}]"
class Matrix:
def __init__(self, inpt: str) -> None:
lines = inpt.splitlines()
self.height = len(lines)
self.width = len(lines[0])
self.array = array("u")
self.array_rot = array("u")
for line in lines:
self.array.fromlist(list(line))
for line in zip(*lines):
self.array_rot.fromlist(list(line))
def __getitem__(self, pos: Position | tuple[int, int]) -> str:
return self.array[self.width * pos[1] + pos[0]]
def __setitem__(self, pos: Position | tuple[int, int], value: str) -> None:
self.array[self.width * pos[1] + pos[0]] = value
def get_row(self, row_no: int) -> array[str]:
return self.array[self.width * row_no : self.width * (row_no + 1)]
def iter_rows(self) -> Iterable[array[str]]:
return (self.get_row(y) for y in range(self.height))
def get_col(self, col_no: int) -> array[str]:
return self.array_rot[self.height * col_no : self.width * (col_no + 1)]
# return array("u", [self[col_no, y] for y in range(self.height)])
def iter_cols(self) -> Iterable[array[str]]:
return (self.get_col(x) for x in range(self.width))
def __str__(self) -> str:
return "\n".join((x.tounicode() for x in self.iter_rows()))
def finditer(self, element: str) -> Iterable[Position]:
return (
Position(x, y)
for y, line in enumerate(self.iter_rows())
for x, char in enumerate(line)
if char == element
)
def filter_row_ids(self, predicate: Callable[[array[str]], bool]) -> Iterable[int]:
return (y for y, row in enumerate(self.iter_rows()) if predicate(row))
def filter_col_ids(self, predicate: Callable[[array[str]], bool]) -> Iterable[int]:
return (x for x, col in enumerate(self.iter_cols()) if predicate(col))
@dataclass
class Map:
map: Matrix
galaxies: Iterable[tuple[Position, Position]]
@classmethod
def from_input(cls, inpt: str) -> Map:
matrix = Matrix(inpt)
missing_xs = array(
"I",
accumulate([1 if col.count("#") == 0 else 0 for col in matrix.iter_cols()]),
)
missing_ys = array(
"I",
accumulate([1 if row.count("#") == 0 else 0 for row in matrix.iter_rows()]),
)
offset_galaxies = (
(
Position(
missing_xs[pos.x],
missing_ys[pos.y],
),
pos,
)
for pos in matrix.finditer("#")
)
return cls(matrix, offset_galaxies)
class Day11(AocSameParser[Map], Aoc2[Map, Map]):
def parseinput(self, inpt: str) -> Map:
return Map.from_input(inpt)
def part1(self, inpt: Map) -> int:
return sum(
g1 | g2 for g1, g2 in combinations(((o + g) for o, g in inpt.galaxies), 2)
)
def part2(self, inpt: Map) -> int:
return sum(
g1 | g2
for g1, g2 in combinations(((999_999 * o + g) for o, g in inpt.galaxies), 2)
)
def main() -> None:
day11 = Day11(2023, 11)
day11.run()
if __name__ == "__main__":
main()