diff --git a/2023/d14/ex2/ex2.py b/2023/d14/ex2/ex2.py new file mode 100755 index 0000000..090724f --- /dev/null +++ b/2023/d14/ex2/ex2.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python + +import dataclasses +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 tilt_south(self) -> "Map": + rocks = [[Cell.EMPTY for _ in range(self.rows)] for _ in range(self.lines)] + + for y in range(self.rows): + rolling_stop = self.rows + + for x in reversed(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 tilt_west(self) -> "Map": + rocks = [[Cell.EMPTY for _ in range(self.rows)] for _ in range(self.lines)] + + for x in range(self.lines): + rolling_stop = -1 + + for y in range(self.rows): + # 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 = y + rocks[x][y] = Cell.CUBE + continue + # For rollers, roll it up to the `last_cube` + rocks[x][rolling_stop + 1] = Cell.ROLLER + rolling_stop += 1 + + return Map(tuple(map(tuple, rocks)), self.lines, self.rows) + + def tilt_east(self) -> "Map": + rocks = [[Cell.EMPTY for _ in range(self.rows)] for _ in range(self.lines)] + + for x in range(self.lines): + rolling_stop = self.lines + + for y in reversed(range(self.rows)): + # 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 = y + rocks[x][y] = Cell.CUBE + continue + # For rollers, roll it up to the `last_cube` + rocks[x][rolling_stop - 1] = Cell.ROLLER + rolling_stop -= 1 + + return Map(tuple(map(tuple, rocks)), self.lines, self.rows) + + def cycle(self) -> "Map": + return self.tilt_north().tilt_west().tilt_south().tilt_east() + + 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])) + + def do_cycles(map: Map) -> Map: + cache = {map: 0} + t = 0 + SPIN_CYCLE_LENGTH = 1000000000 + while t < SPIN_CYCLE_LENGTH: + map = map.cycle() + t += 1 + if map in cache: + previous_t = cache[map] + cycle_length = t - previous_t + num_cycles = (SPIN_CYCLE_LENGTH - t) // cycle_length + t += num_cycles * cycle_length + else: + cache[map] = t + + return map + + map = parse(input) + map = do_cycles(map) + return map.load() + + +def main() -> None: + input = sys.stdin.read().splitlines() + print(solve(input)) + + +if __name__ == "__main__": + main()