2022: d17: ex1: add solution

2022/d17/ex1/ex1.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python
+
+import enum
+import itertools
+import sys
+from collections.abc import Iterator
+from typing import NamedTuple
+
+
+class Point(NamedTuple):
+    x: int
+    y: int
+
+    def __add__(self, other):
+        if not isinstance(other, Point):
+            return NotImplemented
+        return Point(self.x + other.x, self.y + other.y)
+
+    def __sub__(self, other):
+        if not isinstance(other, Point):
+            return NotImplemented
+        return Point(self.x - other.x, self.y - other.y)
+
+
+def translate(points: set[Point], delta: Point) -> set[Point]:
+    return {p + delta for p in points}
+
+
+class Rock(str, enum.Enum):
+    LINE = "####"
+    PLUS = ".#.\n###\n.#."
+    CORNER = "..#\n..#\n###"
+    VERTICAL_LINE = "#\n#\n#\n#"
+    SQUARE = "##\n##"
+
+    @classmethod
+    def stream(cls) -> Iterator["Rock"]:
+        yield from itertools.cycle(iter(cls))
+
+    def to_points(self) -> set[Point]:
+        res: set[Point] = set()
+
+        for y, line in enumerate(reversed(self.splitlines())):
+            for x, c in enumerate(line):
+                if c == ".":
+                    continue
+                res.add(Point(x, y))
+
+        return res
+
+
+class JetStream(str, enum.Enum):
+    LEFT = "<"
+    RIGHT = ">"
+
+    @classmethod
+    def stream(cls, jet_pattern: str) -> Iterator["JetStream"]:
+        yield from itertools.cycle(map(cls, jet_pattern))
+
+    def as_delta(self) -> Point:
+        if self == self.LEFT:
+            return Point(-1, 0)
+        if self == self.RIGHT:
+            return Point(1, 0)
+        assert False  # Sanity check
+
+
+def solve(input: list[str]) -> int:
+    assert len(input) == 1  # Sanity check
+
+    rocks = Rock.stream()
+    jet_stream = JetStream.stream(input[0])
+
+    fallen_stack: set[Point] = set()
+    max_height = 0
+
+    LEFT_WALL = -1
+    RIGHT_WALL = 7
+    FLOOR = 0
+
+    def step(rock: set[Point]) -> tuple[set[Point], bool]:
+        jet = next(jet_stream)
+
+        # Check if it can be pushed by the jet, or if it hits an obstacle
+        pushed_rock = translate(rock, jet.as_delta())
+        if not (fallen_stack & pushed_rock) and all(
+            LEFT_WALL < p.x < RIGHT_WALL for p in pushed_rock
+        ):
+            rock = pushed_rock
+
+        # Check if it can go down
+        fallen_rock = translate(rock, Point(0, -1))
+
+        if not (fallen_stack & fallen_rock) and all(p.y > FLOOR for p in fallen_rock):
+            return fallen_rock, True
+        return rock, False
+
+    def simulate_rock_fall() -> None:
+        nonlocal max_height
+
+        rock = next(rocks).to_points()
+
+        # Align 2 units away from LEFT_WALL and 3 higher than
+        # current stack
+        rock = translate(rock, Point(2, max_height + 3 + 1))
+
+        while True:
+            rock, keep_going = step(rock)
+            if not keep_going:
+                break
+
+        fallen_stack.update(rock)
+        max_height = max(max_height, max(p.y for p in rock))
+
+    for _ in range(2022):
+        simulate_rock_fall()
+
+    return max_height
+
+
+def main() -> None:
+    input = sys.stdin.read().splitlines()
+    print(solve(input))
+
+
+if __name__ == "__main__":
+    main()