advent-of-code/2020/day11/day11.py

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