advent-of-code/2020/solutions/11/python/main.py

from os.path import join, dirname
from typing import List, Optional
from itertools import product
from copy import deepcopy

TSeats = List[List[Optional[bool]]]

mapping = {
    '.': None,
    '#': True,
    'L': False
}
inv = {v: k for k, v in mapping.items()}


class Seats:
    p = [-1, 0, 1]

    def __init__(self, plan: str, alt: bool = False) -> None:
        self.seats: TSeats = [
            [
                mapping[seat]
                for seat in row
            ]
            for row in plan.strip().split('\n')
        ]
        self.max_x = len(self.seats[0])
        self.max_y = len(self.seats)
        self.alt = alt

    def __str__(self) -> str:
        return '\n'.join([
            ''.join([inv[seat] for seat in row])
            for row in self.seats
        ])

    def find_next_in_direction(self, y: int, x: int, dy: int, dx: int) -> Optional[bool]:
        y += dy
        x += dx
        while 0 <= x < self.max_x and 0 <= y < self.max_y:
            cur = self.seats[y][x]
            if cur is not None:
                return cur
            y += dy
            x += dx
        return None

    def get_occupied(self, y: int, x: int,) -> int:
        occupied = 0
        for dx, dy in product(self.p, self.p):
            if dx == 0 and dy == 0:
                continue
            if self.alt and self.find_next_in_direction(y, x, dy, dx) == True:
                occupied += 1
            else:
                dx += x
                dy += y
                if 0 <= dx < self.max_x and 0 <= dy < self.max_y and self.seats[dy][dx]:
                    occupied += 1
        return occupied

    def iteration(self) -> int:
        changed = 0
        future: TSeats = deepcopy(self.seats)
        required_to_leave = 4 if self.alt else 3
        for y, x in product(range(self.max_y), range(self.max_x)):
            current = self.seats[y][x]
            if current == None:
                continue
            occupied = self.get_occupied(y, x)
            if (current == True and occupied > required_to_leave) or (current == False and occupied == 0):
                future[y][x] = not current
                changed += 1
        self.seats = future
        return changed

    def count_occupied(self) -> int:
        return sum([
            sum([
                1 if seat == True else 0
                for seat in row
            ])
            for row in self.seats
        ])

    def find_equilibrium(self) -> int:
        while self.iteration() > 0:
            pass
        return self.count_occupied()


data = join(dirname(__file__), '../data.txt')
with open(data) as f:
    txt = f.read()
    seats = Seats(txt)
    print(seats.find_equilibrium())

    seats = Seats(txt, True)
    print(seats.find_equilibrium())