#!/usr/bin/env python from copy import deepcopy from enum import Enum, auto from functools import cache from itertools import chain import sys from typing import Iterable, List, Optional, Tuple NEIGHBOURS = [(dx, dy) for dy in (-1, 0, 1) for dx in (-1, 0, 1) if not dx == dy == 0] class Tile(Enum): FLOOR = '.' FREE = 'L' OCCUPIED = '#' class Layout: def __init__( self, s: str, max_neighbour_steps: Optional[int], occupied_neighbours_to_leave: int): self._store = [[Tile(c) for c in line] for line in s.splitlines()] self.num_rows = len(self._store) self.num_cols = len(self._store[0]) self.max_neighbour_steps = max_neighbour_steps self.occupied_neighbours_to_leave = occupied_neighbours_to_leave def tiles(self) -> Iterable[Tile]: return chain.from_iterable(self._store) def get(self, row: int, col: int) -> Tile: return self._store[row][col] def set(self, row: int, col: int, t: Tile) -> None: self._store[row][col] = t def step(self) -> bool: new_tiles = [] for row in range(self.num_rows): for col in range(self.num_cols): num_neighbours = self.count_neighbours(row, col) tile = self.get(row, col) new_tile = None if tile == Tile.FREE and num_neighbours == 0: new_tile = Tile.OCCUPIED elif tile == Tile.OCCUPIED and num_neighbours >= self.occupied_neighbours_to_leave: new_tile = Tile.FREE if new_tile is not None: new_tiles.append((row, col, new_tile)) for row, col, tile in new_tiles: self.set(row, col, tile) return len(new_tiles) != 0 def _find_neighbour(self, row: int, col: int, direction: Tuple[int, int]) -> Optional[Tuple[int, int]]: steps = 0 while self.max_neighbour_steps is None or steps < self.max_neighbour_steps: row += direction[0] col += direction[1] if row < 0 or row >= self.num_rows: break if col < 0 or col >= self.num_cols: break if self.get(row, col) != Tile.FLOOR: return (row, col) steps += 1 return None @cache def _neighbour_coords(self, row: int, col: int) -> List[Tuple[int, int]]: coords = [] for direction in NEIGHBOURS: neighbour = self._find_neighbour(row, col, direction) if neighbour is not None: coords.append(neighbour) return coords def neighbours(self, row: int, col: int) -> Iterable[Tile]: for nrow, ncol in self._neighbour_coords(row, col): yield self.get(nrow, ncol) def count_neighbours(self, row: int, col: int) -> int: return sum(1 for tile in self.neighbours(row, col) if tile == Tile.OCCUPIED) def run(layout: Layout) -> int: while layout.step(): pass return len([t for t in layout.tiles() if t == Tile.OCCUPIED]) if __name__ == '__main__': layout_str = sys.stdin.read() print(run(Layout(layout_str, 1, 4))) print(run(Layout(layout_str, None, 5)))