2023: d10: ex2: add solution

2023/d10/ex2/ex2.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python
+
+import sys
+from collections.abc import Iterator
+from enum import Enum, StrEnum
+from typing import NamedTuple, Optional
+
+
+class Point(NamedTuple):
+    x: int
+    y: int
+
+
+class Direction(Enum):
+    NORTH = Point(-1, 0)
+    SOUTH = Point(1, 0)
+    WEST = Point(0, -1)
+    EAST = Point(0, 1)
+
+    def apply(self, pos: "Point") -> "Point":
+        dx, dy = self.value
+        return Point(pos.x + dx, pos.y + dy)
+
+
+class Pipe(StrEnum):
+    VERTICAL = "|"
+    HORIZONTAL = "-"
+    NE_BEND = "L"
+    NW_BEND = "J"
+    SW_BEND = "7"
+    SE_BEND = "F"
+
+    @classmethod
+    def from_connection(cls, connections: dict[Direction, bool]) -> "Pipe":
+        assert sum(connections.values()) == 2  # Sanity check
+        if connections.get(Direction.NORTH, False):
+            if connections.get(Direction.SOUTH, False):
+                return cls.VERTICAL
+            if connections.get(Direction.EAST, False):
+                return cls.NE_BEND
+            if connections.get(Direction.WEST, False):
+                return cls.NW_BEND
+        if connections.get(Direction.SOUTH, False):
+            if connections.get(Direction.WEST, False):
+                return cls.SW_BEND
+            if connections.get(Direction.EAST, False):
+                return cls.SE_BEND
+        assert connections[Direction.WEST] and connections[Direction.EAST]
+        return cls.HORIZONTAL
+
+    def neighbours(self, pos: Point) -> Iterator[Point]:
+        deltas: tuple[Direction, Direction]
+
+        match self:
+            case self.VERTICAL:
+                deltas = (Direction.NORTH, Direction.SOUTH)
+            case self.HORIZONTAL:
+                deltas = (Direction.WEST, Direction.EAST)
+            case self.NE_BEND:
+                deltas = (Direction.NORTH, Direction.EAST)
+            case self.NW_BEND:
+                deltas = (Direction.NORTH, Direction.WEST)
+            case self.SW_BEND:
+                deltas = (Direction.SOUTH, Direction.WEST)
+            case self.SE_BEND:
+                deltas = (Direction.SOUTH, Direction.EAST)
+
+        for dir in deltas:
+            yield dir.apply(pos)
+
+    def go_through(self, pos: Point, prev: Point) -> Point:
+        for dest in self.neighbours(pos):
+            if dest == prev:
+                continue
+            return dest
+        assert False  # Sanity check
+
+
+Pipes = dict[Point, Pipe]
+
+
+def solve(input: list[str]) -> int:
+    def parse(input: list[str]) -> tuple[Point, Pipes]:
+        start: Optional[Point] = None
+        pipes: Pipes = {}
+
+        for x, line in enumerate(input):
+            for y, c in enumerate(line):
+                if c == ".":
+                    continue
+                if c == "S":
+                    start = Point(x, y)
+                    continue
+                pipes[Point(x, y)] = Pipe(c)
+
+        assert start is not None  # Sanity check
+        return start, pipes
+
+    def iter_pipe(cur: Point, start: Point, pipes: Pipes) -> Iterator[Point]:
+        prev = start
+        while True:
+            if cur == start:
+                break
+            if cur not in pipes:
+                break
+            # Check that the receiving pipe is connected to previous one
+            if prev not in pipes[cur].neighbours(cur):
+                return
+            yield cur
+            cur, prev = pipes[cur].go_through(cur, prev), cur
+        # Yield the loop element
+        if cur == start:
+            yield cur
+
+    # Return the possible pipes that start can be
+    def resolve_start(start: Point, pipes: Pipes) -> Pipe:
+        # Returns the direction which closes the loop, if there is one
+        def explore(dir: Direction) -> Optional[Direction]:
+            points = list(iter_pipe(dir.apply(start), start, pipes))
+            if len(points) == 0:
+                return None
+            penultimate, last = points[-2], points[-1]
+            if last != start:
+                return None
+            return Direction(Point(penultimate.x - last.x, penultimate.y - last.y))
+
+        res: set[Pipe] = set()
+        for dir in Direction:
+            if (resulting_dir := explore(dir)) is None:
+                continue
+            res.add(Pipe.from_connection({dir: True, resulting_dir: True}))
+        assert len(res) == 1  # Instructions say there is exactly one loop
+        return res.pop()
+
+    def loop_points(start: Point, pipes: Pipes) -> list[Point]:
+        return list(iter_pipe(next(pipes[start].neighbours(start)), start, pipes))
+
+    def shoelace_area(points: list[Point]) -> int:
+        # Must be integer because pipes follow the grid, and can't cut squares in half
+        return (
+            sum(
+                (points[i - 1].x * points[i].y) - (points[i].x * points[i - 1].y)
+                for i in range(len(points))
+            )
+            // 2
+        )
+
+    start, pipes = parse(input)
+    pipes[start] = resolve_start(start, pipes)
+    loop = loop_points(start, pipes)
+    area = abs(shoelace_area(loop))
+    # len(loop) must be pair, because it's a *loop*
+    return area - len(loop) // 2 + 1  # Using Pick's formula
+
+
+def main() -> None:
+    input = sys.stdin.read().splitlines()
+    print(solve(input))
+
+
+if __name__ == "__main__":
+    main()