Let's play some Tetris with volcanic rock!
From the puzzle description we learn:
The floor is y coordinate 0, the left wall makes x 0, and for each rock shape, we can create a set-like object that contains the coordinates of the rock shapes. The actual shapes are just the relative points from its bottom left corner.
from dataclasses import dataclass
from enum import Enum
from itertools import cycle, islice
from typing import Final, Iterator, NamedTuple, Self
WIDTH: Final[int] = 7
class Pos(NamedTuple):
x: int
y: int
def __add__(self, other: Self) -> Self:
if not isinstance(other, Pos):
return NotImplemented
return Pos(self.x + other.x, self.y + other.y)
def __sub__(self, other: Self) -> Self:
if not isinstance(other, Pos):
return NotImplemented
return Pos(self.x - other.x, self.y - other.y)
class Jet(Enum):
left = ("<", -1)
right = (">", 1)
dx: int
def __new__(cls, char: str, dx: int) -> Self:
instance = object.__new__(cls)
instance._value_ = char
instance.dx = dx
return instance
def __radd__(self, pos: Pos) -> Pos:
return Pos(pos.x + self.dx, pos.y)
@dataclass(frozen=True)
class RockShape:
pos: Pos
width: int
height: int
points: frozenset[Pos]
@classmethod
def from_text(cls, text: str) -> Self:
lines = text.splitlines()
points = frozenset(
Pos(x, y)
for y, line in enumerate(reversed(lines))
for x, char in enumerate(line)
if char == "#"
)
return cls(Pos(0, 0), len(lines[0]), len(lines), points)
def __add__(self, v: Pos | Jet) -> Self:
return __class__(self.pos + v, self.width, self.height, self.points)
def __iter__(self) -> Iterator[Pos]:
pos = self.pos
yield from (pos + p for p in self.points)
def isdisjoint(self, other: set[Pos]) -> bool:
return not any(p in other for p in self)
@property
def x(self) -> int:
return self.pos.x
@property
def y(self) -> int:
return self.pos.y
ROCKS: Final[tuple[RockShape, ...]] = (
RockShape.from_text("####"),
RockShape.from_text(".#.\n###\n.#."),
RockShape.from_text("..#\n..#\n###"),
RockShape.from_text("#\n#\n#\n#"),
RockShape.from_text("##\n##"),
)
class CycleState(NamedTuple):
height: int
period: int
rock_delta: Pos
class VolcanicRockFall:
jets: list[Jet]
def __init__(self, jets: str) -> None:
self.jets = [Jet(j) for j in jets]
def simulate(self) -> Iterator[CycleState]:
height = jrindex = 0
jets, rocks = cycle(self.jets), cycle(ROCKS)
jrindex = cycle(range(len(ROCKS) * len(self.jets)))
settled: set[Pos] = set()
for rock in rocks:
rock += Pos(2, height + 3)
last_jri = 0
for jet, jri in zip(jets, jrindex):
last_jri = jri
# push sideways and test against walls & settled rocks
pushed = rock + jet
if 0 <= pushed.x <= WIDTH - pushed.width and (
pushed.y > height or pushed.isdisjoint(settled)
):
rock = pushed
# fall one step down, see if it settled
fallen = rock + Pos(0, -1)
if fallen.y < 0 or (
fallen.y < height and not fallen.isdisjoint(settled)
):
break
rock = fallen
settled.update(rock)
rock_delta = Pos(2, height + 3) - rock.pos
height = max(rock.height + rock.y, height)
yield CycleState(height, last_jri, rock_delta)
def height_after(self, cycles: int) -> int:
# skip cycles - 1 steps, next value is our step to take info for.
return next(islice(self.simulate(), cycles - 1, None)).height
example = VolcanicRockFall(">>><<><>><<<>><>>><<<>>><<<><<<>><>><<>>")
assert example.height_after(2022) == 3068
import aocd
rockfall = VolcanicRockFall(aocd.get_data(day=17, year=2022))
print("Part 1:", rockfall.height_after(2022))
Part 1: 3209
No implementation is going to be efficient enough to run the simulation a trillion steps, no amount of optimisation is going to make this fast.
The trick is to look for cycles. The number of jet movements times the number of rocks forms a repeating pattern of falling rocks, but because of the possibly endless gaps between settled rocks it'll take a while for a pattern to actually emerge, because the number of jet movements used for each falling rock will vary. To detect when a cycle occurs, you need to track a key indicating the current state, and a series of step numbers and heights at which we have seen that state. Once the delta between the height now and the previous time we saw this state is the same as the height between the state then and the state before it, you know you can use that height difference to cover for most of the rest of the cycles (depending on the modulo of the remainder and the cycle length).
For the key, we can use the rock-and-jet index (a number between 0 and number-of-rocks times number-of-jets) at the point where the rock settled, and the delta of the rock position between where it entered and where it settled (to account for the gradual changes in the stacked stone). The value is a deque with a maximum length of 2 so older values are evicted automatically, storing the step number and the height at reached at that step.
I refactored my implementation to just yield the height and the state key each time the rock settled, then implemented the history tracking based on this.
from collections import defaultdict, deque
from contextlib import suppress
from functools import partial
ONE_TRILLION: Final[int] = 1_000_000_000_000
class HistoryEntry(NamedTuple):
step: int
height: int
def one_trillionth_height(rockfall: VolcanicRockFall) -> int:
heights_queue_factory = partial(deque, (), 2)
history: dict[CycleState, deque[HistoryEntry]] = defaultdict(heights_queue_factory)
skip_height = skip_steps = 0
for step, state in enumerate(rockfall.simulate()):
if not skip_height:
entries = history[state.period, state.rock_delta]
with suppress(ValueError):
h0, h1 = entries
if (hdelta := h1.height - h0.height) == state.height - h1.height:
# cycle detected, skip ahead
cycle_length = step - h1.step
cycles_to_skip = (ONE_TRILLION - step - 1) // cycle_length
skip_height, skip_steps = (
hdelta * cycles_to_skip,
cycle_length * cycles_to_skip,
)
entries.append(HistoryEntry(step, state.height))
if step + skip_steps == ONE_TRILLION - 1:
return state.height + skip_height
assert one_trillionth_height(example) == 1514285714288
print("Part 2:", one_trillionth_height(rockfall))
Part 2: 1580758017509