100 lines
2.7 KiB
Python
100 lines
2.7 KiB
Python
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())
|