Smoke Basin

So using NumPy is almost starting to feel like cheating… However, it is the best way in Python, and I do maintain (strongly) that it should be considered for the standard library. NumPy is the bedrock upon which Python’s foothold in numerical computation rests. If it was banned, we might as well be using Julia.

Alright, so technically here I’m using SciPy methods, but they’re built on NumPy.

Prelude

Standard fetch & setup.

aoc = __import__('aoc-common').aoc()
import numpy as np

text = aoc.input_text(9)

# Packed single-digit text matrix to numpy.
lines = text.split()
rows = len(lines)
heights = np.frombuffer(''.join(lines).encode(), 'u1').reshape(rows, -1) - ord('0')

Part 1

Here’s another task where it’d be much better if both problems were revealed at once. I started out writing something a bit less elegant here, like adding a border and doing (h[1:-1,:] < h[2:,:]) + (h[1:-1,:] > h[:-2,:]) + ....

But then part 2 revealed they obviously want me to use ndimage.label, so alright, and I went back and fixed part 1 since if I’m using ndimage.label anyway, it can be simplified greatly.

import scipy.ndimage as ndi

# Cycle the digits so 9 becomes 0.
#
# Besides, part 1 asks for +1 on minima, so it's like it was designed for this.
heights = (heights + 1) % 10

# It turns out there's a guarantee of the input that make it so all 'regions'
# will be separated by 9 and only 9 (conveniently 0 in our offset image).
lbls,nlbls = ndi.label(heights)
lix = np.arange(1, nlbls+1)

answer1 = ndi.minimum(heights, lbls, lix).astype('i8').sum()
print(f'answer to part1: {answer1}')

Part 2

I don’t really think there’s much point in doing a conventional breadth-first search, even for illustrative purposes. I’d rather someone learn ndimage.label than write a BFS in pure Python for something like this.

# Unfortunately sum_labels seem to only return floats, which we don't want.
basin_sizes = ndi.labeled_comprehension(heights, lbls, lix, np.count_nonzero, np.int64, -1)

answer2 = np.prod(np.sort(basin_sizes)[-3:])
print(f'answer to part1: {answer2}')