DEV Community

Yuan Gao
Yuan Gao

Posted on • Updated on

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()]
Enter fullscreen mode Exit fullscreen mode

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]]
Enter fullscreen mode Exit fullscreen mode
a = 0 + 1j
b = 1 + 1j
c = a + b
Enter fullscreen mode Exit fullscreen mode

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)
Enter fullscreen mode Exit fullscreen mode

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))
Enter fullscreen mode Exit fullscreen mode

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

multiplied = pos.real * pos.imag
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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]]
Enter fullscreen mode Exit fullscreen mode

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])
Enter fullscreen mode Exit fullscreen mode

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

multiplied = pos[0]*pos[1]
Enter fullscreen mode Exit fullscreen mode

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 )
Enter fullscreen mode Exit fullscreen mode

Discussion (2)

Collapse
yassinealouini profile image
Yassine Alouini

I think there is a typo:

new_depth = depth * aim

should be

new_depth = depth + amount * aim

Collapse
meseta profile image
Yuan Gao Author

good call, thanks I fixed