Advent of Code 2021: Day 01 with Python and numpy's convolution function

Link to challenge on Advent of Code 2021 website

Loading the data

Loading the data is straightforward with numpy, which by default can load new-line separated numbers into a numpy array and do the string-to-number conversion in one go:

import numpy as np

lines = np.loadtxt("day_1.txt")

Part 1

Part 1 involves finding how many numbers in a list are greater than the previous number. With Numpy, this is actually quite simple, because it'll do element-wise calculations on whole lists for you.

So writing lines[1:] > lines[:-1] means "do the greater-than comparison for each element on a list that is a shifted version of the data and itself (minus the last element, to match up the lengths)". The result is an array of True or False depending on whether that position had an increase or not. We can sum those (effectively counting the True values) to find the result:

increases = (lines[1:] > lines[:-1]).sum()

Part 2

In part 2, the question introduces a sliding window sum to smooth out the values. The task is to check if the sliding window sum goes up or down. An easy way to implement this is to recognize that a sliding window is a convolution operation with the kernel [1, 1, 1]. The "valid" argument ensures there's no extra padding either side the result, so it's the same size as the input data.

convolved = np.convolve(lines, [1, 1, 1], "valid")

The result is an array containing the moving window sums, which we can compute the number of increased values like before:

increases = (convolved[1:] > convolved[:-1]).sum()

Part 2 alternative solution

Another way to approach Part 2 is to realize that comparing two sliding windows is the same as comparing the unique values in them. If I were to write out the calculation going on, it would look like:

a[i+1] + a[i+2] + a[i+3] > a[i] + a[i+1] + a[i+2]

It can be seen that a[i+1] and a[i+2] appear on both sides, and therefore cancel out, leaving just a[i+3] > [ai]. Therefore Part 2 can be achieved using the same element-wise calculation as Part 1, by shifting the lists by 3 rather than 1, skipping needing to do any convolution:

increases = (lines[3:] > lines[:-3]).sum()

Full code:

import numpy as np
lines = np.loadtxt("day_1.txt")

increases = (lines[1:] > lines[:-1]).sum()
print("Part 1 result:", increases)

increases = (lines[3:] > lines[:-3]).sum()
print("Part 2 result:", increases)