139 lines
3.6 KiB
Python
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))
|