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

100 lines
2.7 KiB
Python
Raw Normal View History

2020-12-11 21:22:15 +01:00
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())