153 lines
4.6 KiB
Python
Executable file
153 lines
4.6 KiB
Python
Executable file
#!/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()
|