51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
from tools.aoc import AOCDay
|
|
from tools.coordinate import Coordinate
|
|
from tools.grid import Grid
|
|
from typing import Any, List
|
|
|
|
|
|
def getCaveMapFromInput(puzzle_input: List[str]) -> Grid:
|
|
caveMap = Grid(99)
|
|
for y, s in enumerate(puzzle_input):
|
|
for x, v in enumerate(s):
|
|
caveMap.set(Coordinate(x, y), int(v))
|
|
|
|
return caveMap
|
|
|
|
|
|
def getLowPoints(caveMap: Grid) -> List[Coordinate]:
|
|
lowPoints = []
|
|
for point in caveMap.getActiveCells():
|
|
if not sum((caveMap.get(t) <= caveMap.get(point)) for t in point.getNeighbours(includeDiagonal=False)):
|
|
lowPoints.append(point)
|
|
|
|
return lowPoints
|
|
|
|
|
|
def getBasin(caveMap: Grid, start: Coordinate, visited: set) -> set:
|
|
for n in start.getNeighbours(includeDiagonal=False):
|
|
if n in visited:
|
|
continue
|
|
|
|
if caveMap.get(n) < 9:
|
|
visited.add(n)
|
|
visited = set.union(visited, getBasin(caveMap, n, visited))
|
|
|
|
return visited
|
|
|
|
|
|
class Day(AOCDay):
|
|
test_solutions_p1 = [15]
|
|
test_solutions_p2 = [1134]
|
|
|
|
def part1(self) -> Any:
|
|
caveMap = getCaveMapFromInput(self.getInput())
|
|
return sum(caveMap.get(x) + 1 for x in getLowPoints(caveMap))
|
|
|
|
def part2(self) -> Any:
|
|
caveMap = getCaveMapFromInput(self.getInput())
|
|
lowPoints = getLowPoints(caveMap)
|
|
basins = list(sorted([getBasin(caveMap, point, set()) for point in lowPoints], key=lambda l: len(l)))
|
|
|
|
return len(basins[-1]) * len(basins[-2]) * len(basins[-3])
|