aoc-2023/d11/main.py
2023-12-12 14:01:10 -08:00

139 lines
3.6 KiB
Python

from collections.abc import Generator
from itertools import combinations, product
from pathlib import Path
from typing import NamedTuple
class Point(NamedTuple):
x: int
y: int
def iter_adjacent(self, wrap=False) -> Generator["Point", None, None]:
for offset in product((1, 0, -1), (1, 0, -1)):
adj = Point(self.x + offset[0], self.y + offset[1])
if adj == self:
continue
if not wrap and (adj.x < 0 or adj.y < 0):
continue
yield adj
def add_point(self, delta: "Point") -> "Point":
return Point(self.x + delta.x, self.y + delta.y)
def add(self, x=0, y=0) -> "Point":
return Point(self.x + x, self.y + y)
def dist(self, x=0, y=0, point: "Point|None" = None) -> int:
if (x or y) and point:
raise ValueError("Cannot provide both x and y and a point")
if point:
x = point.x
y = point.y
dx = abs(self.x - x)
dy = abs(self.y - y)
return dx + dy
def __str__(self) -> str:
return f"Point({self.x}, {self.y})"
class Galaxy:
def __init__(self, lines: list[str]):
self.planets: set[Point] = set()
self.grid = [[x for x in row] for row in lines if row]
self.expand_y: dict[int, int] = {}
self.expand_x: dict[int, int] = {}
def find_planets(self) -> set[Point]:
# Find planets
for y, row in enumerate(self.grid):
for x, c in enumerate(row):
if c == "#":
self.planets.add(Point(x, y))
return self.planets
def dist(self, a: Point, b: Point) -> int:
d = a.dist(point=b)
for y in range(*sorted((a.y, b.y))):
d += self.expand_y.get(y, 0)
for x in range(*sorted((a.x, b.x))):
d += self.expand_x.get(x, 0)
return d
def virtual_expand(self, n=2) -> None:
# Expand y
self.expand_y = {
y: n-1 for y, row in enumerate(self.grid) if all(c == "." for c in row)
}
# Expand x
self.expand_x = {
x: n-1 for x in range(len(self.grid[0])) if all(row[x] == "." for row in self.grid)
}
def expand(self, n = 2) -> None:
# Expand y
og_y = len(self.grid)
for y in range(og_y):
i = og_y - y - 1
row = self.grid[i]
if all(c == "." for c in row):
self.grid.insert(i, [str(n-1) for _ in row])
# Expand x
og_x = len(self.grid[0])
for x in range(og_x):
i = og_x - x - 1
col = [row[i] for row in self.grid]
if all(c == "." or c.isnumeric() for c in col):
for row in self.grid:
row.insert(i, str(n-1))
def __str__(self) -> str:
return "\n".join("".join(row) for row in self.grid)
def part1(input: Path) -> int:
galaxy = Galaxy(input.read_text().split("\n"))
galaxy.expand()
galaxy.find_planets()
return sum(
a.dist(point=b) for a, b in combinations(galaxy.planets, 2)
)
def part2(input: Path, n=1000000) -> int:
galaxy = Galaxy(input.read_text().split("\n"))
galaxy.virtual_expand(n)
galaxy.find_planets()
# print(galaxy.dist(Point(1, 0), Point(3, 1)))
return sum(
galaxy.dist(a, b) for a, b in combinations(galaxy.planets, 2)
)
if __name__ == "__main__":
# input = Path("sample.txt")
input = Path("input.txt")
result1 = part1(input)
print("part 1", result1)
print("part 2 10", part2(input, 10))
print("part 2 100", part2(input, 100))
print("part 2 1M", part2(input))