Advent of Code 2021: Day 05 with Python, even more numpy

Link to challenge on Advent of Code 2021 website

Loading the data

This time, there's no shortcuts to loading data, it's regex time. The data is a pair of coordinates, and an arrow, and then another pair of coordinates:

578,391 -> 578,322
274,585 -> 651,962
482,348 -> 294,348

We can read this in using regex pattern (\d+),(\d+) -> (\d+),(\d+). We also grab the max bounds of the coordinates as well.

lines = open("day_5.txt").readlines()
coords = np.array([re.match('(\d+),(\d+) -> (\d+),(\d+)', line).groups() for line in lines]).astype(int)
size = np.max(coords)+1

Part 1

Part 1 asks us to draw these lines on the grid, and find out how many lines overlap. It also asks us to only consider horizontal and vertical lines.

Horizontal and vertical lines are characterized by having x or y coordinates equal, so we can filter those out:

hv = coords[(coords[:, 0] == coords[:, 2]) | (coords[:, 1] == coords[:, 3])]

Then for each of these, we increment the specific indices of the array that they cover:

def rrange(start, stop):
    return range(start, stop+1) if stop >= start else range(start, stop-1, -1)

grid = np.zeros((size, size))
for x1, y1, x2, y2 in hv:
    grid[rrange(y1, y2), rrange(x1, x2)] += 1
result = (grid >= 2).sum()

The rrange() function is a small helper to allow range() to run backwards properly. Because range(1, 5) works fine, but range(5, 1) does not unless specifying the step size as -1 with range(5, 1, -1). The coordinates are also adjusted to make the values inclusive.

Once the grid values have been incremented, we can find all the grid elements with values greater or equal to 2, giving the answer.

Part 2

Part 2 asks us to look at not just horizontal and vertical lines, but also diagonal ones. Oddly, we already wrote the code in a way that's general, so no extra code is needed, we just iterate over the original unfiltered coordinates rather than the ones filtered for only horizontal/vertical lines

grid = np.zeros((size, size))
for x1, y1, x2, y2 in coords:
    grid[rrange(y1, y2), rrange(x1, x2)] += 1
result = (grid >= 2).sum()

Full code

import numpy as np
import re

lines = open("day_5.txt").readlines()
coords = np.array([re.match('(\d+),(\d+) -> (\d+),(\d+)', line).groups() for line in lines]).astype(int)
size = np.max(coords)+1

def rrange(start, stop):
    return range(start, stop+1) if stop >= start else range(start, stop-1, -1)

grid = np.zeros((size, size))
hv = coords[(coords[:, 0] == coords[:, 2]) | (coords[:, 1] == coords[:, 3])]
for x1, y1, x2, y2 in hv:
    grid[rrange(y1, y2), rrange(x1, x2)] += 1
result = (grid >= 2).sum()
print("Part 1 result:", result)

grid = np.zeros((size, size))
for x1, y1, x2, y2 in coords:
    grid[rrange(y1, y2), rrange(x1, x2)] += 1
result = (grid >= 2).sum()
print("Part 2 result:", result)