Advent of Code 2021: Day 02 with Python, numpy, and affine transformation matrices

Link to challenge on Advent of Code 2021 website

Loading the data

The input file consists of a direction, and a number on each line, which we can split into an array using line.split(), and iterate over using file.readlines(). Note: in-line opening a file like this with open() outside a context manager is generally bad practice, but relatively benign in a small script like this.

lines = [line.split() for line in open('day_2.txt').readlines()]

Part 1

Part 1 involves interpreting the input file as directions and distances, adding together the distance travelled. If we think about this in mathematical terms, this involves adding together vectors representing depth and forward distance.

When we deal with 2-vectors, Python gives us a little shortcut: native support for complex numbers. We can use a single complex number to represent a 2-vector, and do arithmetic on it as normal. The two are functionally equivalent:

a = [0, 1]
b = [1, 1]
c = [a[0] + b[0], a[1] + b[1]]

a = 0 + 1j
b = 1 + 1j
c = a + b

To start, we create a dictionary of moves. This helps us look up the correct move vector without having to do an if statement

moves = dict(forward=1, down=1j, up=-1j)

Then, we apply (map()) a function that multiplies together one of these move vectors and the distance number provided in the file. And then sum together all the resultant vectors for the final position of the sub.

pos = sum(map(lambda line: moves[line[0]]*int(line[1]), lines))

The result the question asks for is the two position components multiplied together.

multiplied = pos.real * pos.imag

Part 2

Part 2 introduces the idea that there is an "aim" state, and the instructions can change this aim instead of moving the position of the sub. There are many ways to solve this problem, since I already started on using vectors, I've chosen to continue with that in Part 2, treating this as an affine transformation problem, and using numpy dot product to move the sub and its internal "aim" state.

The moves can be expressed as:

# down move
new_aim = aim + amount

# up move
new_aim = aim - amount

# forward move
new_dist = dist + amount
new_depth = depth + depth * aim

In terms of affine transformation matrices, with our augmented vector being [dist, depth, aim, 1] (a is the amount we're moving by)

# down move
[1, 0, 0, 0] # dist
[0, 1, 0, 0] # depth
[0, 0, 1, a] # aim
[0, 0, 0, 1] # 1

# up move
[1, 0, 0, 0] # dist
[0, 1, 0, 0] # depth
[0, 0, 1,-a] # aim
[0, 0, 0, 1] # 1

# forward move
[1, 0, 0, a] # dist
[0, 1, a, 0] # depth
[0, 0, 1, 0] # aim
[0, 0, 0, 1] # 1

To make this a little easier, we can decompose these moves down to the Identity plus the amount multiplied by the non-diagonal matrix: I_4 + amount * M

In code, producing a moves dictionary containing just the non-diagonal unit matrix:

moves = defaultdict(lambda: np.zeros([4, 4]))
moves["down"][2, 3] = 1
moves["up"][2, 3] = -1
moves["forward"][0:2, 2:4] = [[0,1],[1,0]]

Here, defaultdict is used to help us set up the matrices without having to write out every element, since it is sparse.

We can now iterate over each line, fetching the correct transformation matrix, and applying a dot product with the position vector. Using reduce() to avoid writing out a loop.

pos = reduce(lambda pos, line: (np.identity(4)+moves[line[0]]*int(line[1])).dot(pos), lines, [0, 0, 0, 1])

The result, as before, is the first two elements of the position vector multiplied

multiplied = pos[0]*pos[1]

Full Code:

lines = [x.split() for x in open('day_2.txt').readlines()]

moves = dict(forward=1, down=1j, up=-1j)
pos = sum(map(lambda line: moves[line[0]]*int(line[1]), lines))
result = pos.real * pos.imag
print("Part 1 result:", result )

from collections import defaultdict
from functools import reduce
import numpy as np

moves = defaultdict(lambda: np.zeros([4, 4]))
moves["down"][2, 3] = 1
moves["up"][2, 3] = -1
moves["forward"][0:2, 2:4] = [[0,1],[1,0]]

pos = reduce(lambda pos, line: (np.identity(4)+moves[line[0]]*int(line[1])).dot(pos), lines, [0, 0, 0, 1])
result = pos[0]*pos[1]
print("Part 2 result:", result )