69 lines
1.7 KiB
Python
69 lines
1.7 KiB
Python
|
#!/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()
|