aoc_2021_day09.py

import numpy as np
from numpy import ndarray
import matplotlib.pyplot as plt

from functools import reduce


def up(index):
	i, j = index
	return i - 1, j


def down(index):
	i, j = index
	return i + 1, j


def left(index):
	i, j = index
	return i, j - 1


def right(index):
	i, j = index
	return i, j + 1


class HeightMap:
	def __init__(self, data: ndarray):
		self.data = data

	@property
	def shape(self):
		return self.data.shape

	@classmethod
	def from_str(cls, s: str):
		lines = s.strip().splitlines()
		nums = list(map(lambda l: list(map(int, l)), lines))
		return cls(data=np.array(nums))

	def __str__(self):
		return "\n".join(map(lambda row: "".join(map(str, row)), self.data))

	def smaller_than_top(self, index):
		i, j = index
		if i == 0:
			return True
		else:
			return self.data[i, j] < self.data[i - 1, j]

	def smaller_than_bottom(self, index):
		i, j = index
		if i == self.shape[0] - 1:
			return True
		else:
			return self.data[i, j] < self.data[i + 1, j]

	def smaller_than_left(self, index):
		i, j = index
		if j == 0:
			return True
		else:
			return self.data[i, j] < self.data[i, j - 1]

	def smaller_than_right(self, index):
		i, j = index
		if j == self.shape[1] - 1:
			return True
		else:
			return self.data[i, j] < self.data[i, j + 1]

	def is_min(self, index):
		if not len(index) == 2: raise IndexError("Only 2d indices allowed")
		i, j = index

		height = self.data[i, j]

		return self.smaller_than_top(index) and self.smaller_than_bottom(
			index) and self.smaller_than_left(index) and self.smaller_than_right(index)

	def minima(self):
		mins = {}
		for i in range(self.shape[0]):
			for j in range(self.shape[1]):
				if self.is_min((i, j)):
					mins[(i, j)] = self.data[i, j]
		return mins

	def total_risk(self):
		return sum(v + 1 for v in self.minima().values())

	def get(self, index):
		i, j = index
		if i < 0 or j < 0:
			return
		try:
			return self.data[i, j]
		except IndexError:
			return

	def basin(self, index):
		def walk_basin(height_map, index, visited):
			height = height_map.get(index)
			if height is None:
				return
			elif height == 9:
				return
			elif index in visited:
				return
			else:
				visited.add(index)
				walk_basin(height_map, up(index), visited)
				walk_basin(height_map, down(index), visited)
				walk_basin(height_map, left(index), visited)
				walk_basin(height_map, right(index), visited)

		visited = set()
		walk_basin(self, index, visited)
		return visited

	def part2(self):
		sizes = []
		for index, height in self.minima().items():
			sizes.append(len(self.basin(index)))

		sizes.sort()
		return sizes[-1] * sizes[-2] * sizes[-3]


# Tests	
test_input = "2199943210\n3987894921\n9856789892\n8767896789\n9899965678"

hm = HeightMap.from_str(test_input)
print(hm)

mins = hm.minima()
print(mins)

print(hm.total_risk())

print("Part 2: ", hm.part2())

# Part 1
with open("./day09.txt") as f:
	s = f.read()
	hm = HeightMap.from_str(s)
	print(hm.total_risk())

# Part 2
print(hm.part2())