advent-of-code/2023/d14/ex2/ex2.py
Bruno BELANYI f9fc9fbd6b treewide: fix 'ruff check' errors
This is mostly about unused imports.

A couple errors remain, but are fine in my book (using `l` as a variable
name, assigning a lambda to a variable).
2024-11-23 19:24:55 +00:00

152 lines
4.6 KiB
Python
Executable file

#!/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 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()