advent-of-code/2021/09/python/main.py
Niccolo Borgioli 7c399251d6 09
2021-12-09 15:35:29 +01:00

154 lines
4.0 KiB
Python

#!/usr/bin/env python
from os.path import join, dirname
from typing import List
# Day 09
# Common
def read_input(filename):
data = join(dirname(__file__), '..', filename)
with open(data) as f:
return f.read().strip()
test = read_input('test.txt')
data = read_input('input.txt')
def make_string_blue(string):
return f'\033[94m{string}\033[0m'
def make_string_red(string):
return f'\033[91m{string}\033[0m'
class Point:
def __init__(self, x, y):
self.x = x
self.y = y
def __repr__(self):
return f'({self.x}, {self.y})'
def __hash__(self) -> int:
return hash((self.x, self.y))
def __eq__(self, o: object) -> bool:
if not isinstance(o, Point):
return False
return self.x == o.x and self.y == o.y
def value(self, data: List[List[int]]):
return data[self.y][self.x]
class Floor:
def __init__(self, data: List[List[int]]):
self.data = data
self.max_x = len(data[0])
self.max_y = len(data)
def paint(self, blue: List[Point] = [], red: List[Point] = []) -> str:
out = ''
for y in range(self.max_y):
for x in range(self.max_x):
v = str(self.data[y][x])
if Point(x, y) in red:
v = f'\033[31m{v}\033[0m'
elif Point(x, y) in blue:
v = f'\033[34m{v}\033[0m'
out += v
out += '\n'
return out
def get_neighbors(self, point: Point) -> List[Point]:
neighbors = []
if point.x - 1 >= 0:
neighbors.append(Point(point.x - 1, point.y))
if point.x + 1 < self.max_x:
neighbors.append(Point(point.x + 1, point.y))
if point.y - 1 >= 0:
neighbors.append(Point(point.x, point.y - 1))
if point.y + 1 < self.max_y:
neighbors.append(Point(point.x, point.y + 1))
return neighbors
def find_lowest_points(self) -> List[Point]:
lowest_points = []
for y in range(self.max_y):
for x in range(self.max_x):
p = Point(x, y)
neighbors = self.get_neighbors(p)
values = [p.value(self.data) for p in neighbors]
if p.value(self.data) < min(values):
lowest_points.append(p)
return lowest_points
def find_basins(self):
lowest = self.find_lowest_points()
basins = []
for point in lowest:
basin = set([point])
edges = set([point])
while len(edges) > 0:
new_edges = set()
for edge in edges:
neighbors = self.get_neighbors(edge)
for neighbor in neighbors:
if neighbor not in basin and neighbor.value(self.data) < 9:
basin.add(neighbor)
new_edges.add(neighbor)
edges = new_edges
basins.append(basin)
print(self.paint(
[point for basin in basins for point in basin], lowest))
return basins
def flag2(self) -> int:
basins = self.find_basins()
sizes = [len(b) for b in basins]
sizes = sorted(sizes, reverse=True)
total = 1
for size in sizes[0:3]:
total *= size
return total
def flag(self) -> int:
points = self.find_lowest_points()
return sum([
p.value(self.data) + 1
for p in points
])
@staticmethod
def parse(data: str):
return Floor([
list(map(int, list(row)))
for row in data.split('\n')
])
# 1
print('1.')
floor = Floor.parse(test)
print(floor.paint(floor.find_lowest_points()))
print(f'Test: {floor.flag()}')
floor = Floor.parse(data)
print(floor.paint(floor.find_lowest_points()))
print(f'Real: {floor.flag()}')
# 2
print('\n2.')
floor = Floor.parse(test)
print(f'Test: {floor.flag2()}')
floor = Floor.parse(data)
print(f'Real: {floor.flag2()}')