2023: d14: ex1: add solution

2023/d14/ex1/ex1.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python
+
+import itertools
+import sys
+from enum import StrEnum
+from typing import NamedTuple, Optional
+
+
+class Cell(StrEnum):
+    ROLLER = "O"
+    CUBE = "#"
+    EMPTY = "."
+
+
+class Map(NamedTuple):
+    rocks: tuple[tuple[Optional[Cell], ...], ...]
+    lines: int
+    rows: int
+
+    def tilt_north(self) -> "Map":
+        rocks = [[Cell.EMPTY for _ in range(self.rows)] for _ in range(self.lines)]
+
+        for y in range(self.rows):
+            rolling_stop = -1
+
+            for x in range(self.lines):
+                # Nothing to do on empty cell
+                if self.rocks[x][y] == Cell.EMPTY:
+                    continue
+                # Record the new stop point on cubes
+                if self.rocks[x][y] == Cell.CUBE:
+                    rolling_stop = x
+                    rocks[x][y] = Cell.CUBE
+                    continue
+                # For rollers, roll it up to the `last_cube`
+                rocks[rolling_stop + 1][y] = Cell.ROLLER
+                rolling_stop += 1
+
+        return Map(tuple(map(tuple, rocks)), self.lines, self.rows)
+
+    def load(self) -> int:
+        res = 0
+
+        for x, y in itertools.product(range(self.lines), range(self.rows)):
+            if self.rocks[x][y] != Cell.ROLLER:
+                continue
+            res += self.lines - x
+
+        return res
+
+
+def solve(input: list[str]) -> int:
+    def parse(input: list[str]) -> Map:
+        rocks = tuple(tuple(Cell(c) for c in line) for line in input)
+        return Map(rocks, len(input), len(input[0]))
+
+    map = parse(input)
+    map = map.tilt_north()
+    return map.load()
+
+
+def main() -> None:
+    input = sys.stdin.read().splitlines()
+    print(solve(input))
+
+
+if __name__ == "__main__":
+    main()